@Nullable
@Override
public PyType getReturnType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final PyType returnType = typeProvider.getReturnType(this, context);
if (returnType != null) {
returnType.assertValid(typeProvider.toString());
return returnType;
}
}
if (context.maySwitchToAST(this)
&& LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON30)) {
PyAnnotation anno = getAnnotation();
if (anno != null) {
PyClass pyClass = anno.resolveToClass();
if (pyClass != null) {
return new PyClassTypeImpl(pyClass, false);
}
}
}
final PyType docStringType = getReturnTypeFromDocString();
if (docStringType != null) {
docStringType.assertValid("from docstring");
return docStringType;
}
if (context.allowReturnTypes(this)) {
final Ref<? extends PyType> yieldTypeRef = getYieldStatementType(context);
if (yieldTypeRef != null) {
return yieldTypeRef.get();
}
return getReturnStatementType(context);
}
return null;
}
@Nullable
private PyType getGenericReturnType(
@NotNull TypeEvalContext typeEvalContext, @Nullable PyQualifiedExpression callSite) {
if (typeEvalContext.maySwitchToAST(this)
&& LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON30)) {
PyAnnotation anno = getAnnotation();
if (anno != null) {
PyClass pyClass = anno.resolveToClass();
if (pyClass != null) {
return new PyClassTypeImpl(pyClass, false);
}
}
}
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final PyType returnType = typeProvider.getReturnType(this, callSite, typeEvalContext);
if (returnType != null) {
returnType.assertValid(typeProvider.toString());
return returnType;
}
}
final PyType docStringType = getReturnTypeFromDocString();
if (docStringType != null) {
docStringType.assertValid("from docstring");
return docStringType;
}
if (typeEvalContext.allowReturnTypes(this)) {
final PyType yieldType = getYieldStatementType(typeEvalContext);
if (yieldType != null) {
return yieldType;
}
return getReturnStatementType(typeEvalContext);
}
return null;
}
@Nullable
private PyType getReturnType(@NotNull TypeEvalContext context) {
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final Ref<PyType> returnTypeRef = typeProvider.getReturnType(this, context);
if (returnTypeRef != null) {
final PyType returnType = returnTypeRef.get();
if (returnType != null) {
returnType.assertValid(typeProvider.toString());
}
return returnType;
}
}
final PyType docStringType = getReturnTypeFromDocString();
if (docStringType != null) {
docStringType.assertValid("from docstring");
return docStringType;
}
if (context.allowReturnTypes(this)) {
final Ref<? extends PyType> yieldTypeRef = getYieldStatementType(context);
if (yieldTypeRef != null) {
return yieldTypeRef.get();
}
return getReturnStatementType(context);
}
return null;
}
@Nullable
@Override
public PyType getCallType(
@Nullable PyExpression receiver,
@NotNull Map<PyExpression, PyNamedParameter> parameters,
@NotNull TypeEvalContext context) {
return analyzeCallType(context.getReturnType(this), receiver, parameters, context);
}
private static boolean isDynamicallyEvaluated(
@NotNull Collection<PyNamedParameter> parameters, @NotNull TypeEvalContext context) {
for (PyNamedParameter parameter : parameters) {
final PyType type = context.getType(parameter);
if (type instanceof PyDynamicallyEvaluatedType) {
return true;
}
}
return false;
}
@Override
public void visitPyReturnStatement(PyReturnStatement node) {
if (PsiTreeUtil.getParentOfType(node, ScopeOwner.class, true) == myFunction) {
final PyExpression expr = node.getExpression();
PyType returnType;
returnType = expr == null ? PyNoneType.INSTANCE : myContext.getType(expr);
if (!myHasReturns) {
myResult = returnType;
myHasReturns = true;
} else {
myResult = PyUnionType.union(myResult, returnType);
}
}
}
@Nullable
private PyType getReturnType(@NotNull TypeEvalContext context) {
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final Ref<PyType> returnTypeRef = typeProvider.getReturnType(this, context);
if (returnTypeRef != null) {
return derefType(returnTypeRef, typeProvider);
}
}
if (context.allowReturnTypes(this)) {
final Ref<? extends PyType> yieldTypeRef = getYieldStatementType(context);
if (yieldTypeRef != null) {
return yieldTypeRef.get();
}
return getReturnStatementType(context);
}
return null;
}
@Nullable
@Override
public PyType getCallType(
@NotNull TypeEvalContext context, @NotNull PyQualifiedExpression callSite) {
PyType type = null;
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
type = typeProvider.getCallType(this, callSite, context);
if (type != null) {
type.assertValid(typeProvider.toString());
break;
}
}
if (type == null) {
type = context.getReturnType(this);
}
final PyTypeChecker.AnalyzeCallResults results =
PyTypeChecker.analyzeCallSite(callSite, context);
if (results != null) {
return analyzeCallType(type, results.getReceiver(), results.getArguments(), context);
}
return type;
}
@Nullable
private PyType replaceSelf(
@Nullable PyType returnType,
@Nullable PyExpression receiver,
@NotNull TypeEvalContext context) {
if (receiver != null) {
// TODO: Currently we substitute only simple subclass types, but we could handle union and
// collection types as well
if (returnType instanceof PyClassType) {
final PyClassType returnClassType = (PyClassType) returnType;
if (returnClassType.getPyClass() == getContainingClass()) {
final PyType receiverType = context.getType(receiver);
if (receiverType instanceof PyClassType
&& PyTypeChecker.match(returnType, receiverType, context)) {
return returnClassType.isDefinition()
? receiverType
: ((PyClassType) receiverType).toInstance();
}
}
}
}
return returnType;
}
