@Nullable
@Override
public PyType getReturnType(
@NotNull TypeEvalContext context, @Nullable PyQualifiedExpression callSite) {
final PyType type = getGenericReturnType(context, callSite);
if (callSite == null) {
return type;
}
final PyTypeChecker.AnalyzeCallResults results =
PyTypeChecker.analyzeCallSite(callSite, context);
if (PyTypeChecker.hasGenerics(type, context)) {
if (results != null) {
final Map<PyGenericType, PyType> substitutions =
PyTypeChecker.unifyGenericCall(
this, results.getReceiver(), results.getArguments(), context);
if (substitutions != null) {
return PyTypeChecker.substitute(type, substitutions, context);
}
}
return null;
}
if (results != null && isDynamicallyEvaluated(results.getArguments().values(), context)) {
return PyUnionType.createWeakType(type);
} else {
return type;
}
}
@Override
public PyType getReferenceType(
@NotNull PsiElement referenceTarget, TypeEvalContext context, @Nullable PsiElement anchor) {
if (referenceTarget instanceof PyTargetExpression) {
final PyTargetExpression target = (PyTargetExpression) referenceTarget;
if (context.maySwitchToAST(target)) {
// XXX: Requires switching from stub to AST
final PyAnnotation annotation = target.getAnnotation();
if (annotation != null) {
final PyExpression value = annotation.getValue();
if (value != null) {
return getType(value, new Context(context));
}
return null;
}
}
final String comment = target.getTypeCommentAnnotation();
if (comment != null) {
final PyType type = getStringBasedType(comment, referenceTarget, new Context(context));
if (type instanceof PyTupleType) {
final PyTupleExpression tupleExpr =
PsiTreeUtil.getParentOfType(target, PyTupleExpression.class);
if (tupleExpr != null) {
return PyTypeChecker.getTargetTypeFromTupleAssignment(
target, tupleExpr, (PyTupleType) type);
}
}
return type;
}
}
return null;
}
@Nullable
/** Suits when there is no call site(e.g. implicit __iter__ call in statement for) */
public PyType getReturnTypeWithoutCallSite(
@NotNull TypeEvalContext context, @Nullable PyExpression receiver) {
final PyType type = getGenericReturnType(context, null);
if (PyTypeChecker.hasGenerics(type, context)) {
final Map<PyGenericType, PyType> substitutions =
PyTypeChecker.unifyGenericCall(
this, receiver, Maps.<PyExpression, PyNamedParameter>newHashMap(), context);
if (substitutions != null) {
return PyTypeChecker.substitute(type, substitutions, context);
}
return null;
}
return type;
}
@Nullable
@Override
public PyType getCallType(
@NotNull TypeEvalContext context, @NotNull PyCallSiteExpression callSite) {
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final PyType type = typeProvider.getCallType(this, callSite, context);
if (type != null) {
type.assertValid(typeProvider.toString());
return type;
}
}
final PyExpression receiver = PyTypeChecker.getReceiver(callSite, this);
final List<PyExpression> arguments = PyTypeChecker.getArguments(callSite, this);
final List<PyParameter> parameters = PyUtil.getParameters(this, context);
final PyResolveContext resolveContext =
PyResolveContext.noImplicits().withTypeEvalContext(context);
final List<PyParameter> explicitParameters =
PyTypeChecker.filterExplicitParameters(parameters, this, callSite, resolveContext);
final Map<PyExpression, PyNamedParameter> mapping =
PyCallExpressionHelper.mapArguments(arguments, explicitParameters);
return getCallType(receiver, mapping, context);
}
@Nullable
private PyType analyzeCallType(
@Nullable PyType type,
@Nullable PyExpression receiver,
@NotNull Map<PyExpression, PyNamedParameter> parameters,
@NotNull TypeEvalContext context) {
if (PyTypeChecker.hasGenerics(type, context)) {
final Map<PyGenericType, PyType> substitutions =
PyTypeChecker.unifyGenericCall(receiver, parameters, context);
if (substitutions != null) {
type = PyTypeChecker.substitute(type, substitutions, context);
} else {
type = null;
}
}
if (receiver != null) {
type = replaceSelf(type, receiver, context);
}
if (type != null && isDynamicallyEvaluated(parameters.values(), context)) {
type = PyUnionType.createWeakType(type);
}
return type;
}
public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
if (isOperator("and") || isOperator("or")) {
final PyExpression left = getLeftExpression();
final PyType leftType = left != null ? context.getType(left) : null;
final PyExpression right = getRightExpression();
final PyType rightType = right != null ? context.getType(right) : null;
if (leftType == null && rightType == null) {
return null;
}
return PyUnionType.union(leftType, rightType);
}
final PyTypeChecker.AnalyzeCallResults results = PyTypeChecker.analyzeCall(this, context);
if (results != null) {
final PyType type = results.getCallable().getCallType(context, this);
if (!PyTypeChecker.isUnknown(type) && !(type instanceof PyNoneType)) {
return type;
}
}
if (PyNames.COMPARISON_OPERATORS.contains(getReferencedName())) {
return PyBuiltinCache.getInstance(this).getBoolType();
}
return null;
}
@Nullable
@Override
public PyType getCallType(
@NotNull TypeEvalContext context, @NotNull PyCallSiteExpression callSite) {
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final Ref<PyType> typeRef = typeProvider.getCallType(this, callSite, context);
if (typeRef != null) {
return derefType(typeRef, typeProvider);
}
}
final PyExpression receiver = PyTypeChecker.getReceiver(callSite, this);
final Map<PyExpression, PyNamedParameter> mapping =
PyCallExpressionHelper.mapArguments(callSite, this, context);
return getCallType(receiver, mapping, context);
}
@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;
}
