@Nullable
public static PyType getType(@NotNull PsiElement resolved, @NotNull List<PyType> elementTypes) {
final String qualifiedName = getQualifiedName(resolved);
final List<Integer> paramListTypePositions = new ArrayList<>();
final List<Integer> ellipsisTypePositions = new ArrayList<>();
for (int i = 0; i < elementTypes.size(); i++) {
final PyType type = elementTypes.get(i);
if (type instanceof PyTypeParser.ParameterListType) {
paramListTypePositions.add(i);
} else if (type instanceof PyTypeParser.EllipsisType) {
ellipsisTypePositions.add(i);
}
}
if (!paramListTypePositions.isEmpty()) {
if (!("typing.Callable".equals(qualifiedName) && paramListTypePositions.equals(list(0)))) {
return null;
}
}
if (!ellipsisTypePositions.isEmpty()) {
if (!("typing.Callable".equals(qualifiedName) && ellipsisTypePositions.equals(list(0))
|| "typing.Tuple".equals(qualifiedName)
&& ellipsisTypePositions.equals(list(1))
&& elementTypes.size() == 2)) {
return null;
}
}
if ("typing.Union".equals(qualifiedName)) {
return PyUnionType.union(elementTypes);
}
if ("typing.Optional".equals(qualifiedName) && elementTypes.size() == 1) {
return PyUnionType.union(elementTypes.get(0), PyNoneType.INSTANCE);
}
if ("typing.Callable".equals(qualifiedName) && elementTypes.size() == 2) {
final PyTypeParser.ParameterListType paramList =
as(elementTypes.get(0), PyTypeParser.ParameterListType.class);
if (paramList != null) {
return new PyCallableTypeImpl(paramList.getCallableParameters(), elementTypes.get(1));
}
if (elementTypes.get(0) instanceof PyTypeParser.EllipsisType) {
return new PyCallableTypeImpl(null, elementTypes.get(1));
}
}
if ("typing.Tuple".equals(qualifiedName)) {
if (elementTypes.size() > 1 && elementTypes.get(1) instanceof PyTypeParser.EllipsisType) {
return PyTupleType.createHomogeneous(resolved, elementTypes.get(0));
}
return PyTupleType.create(resolved, elementTypes);
}
final PyType builtinCollection = getBuiltinCollection(resolved);
if (builtinCollection instanceof PyClassType) {
final PyClassType classType = (PyClassType) builtinCollection;
return new PyCollectionTypeImpl(classType.getPyClass(), false, elementTypes);
}
return null;
}
@Nullable
private static PyType getParameterizedType(
@NotNull PsiElement element, @NotNull Context context) {
if (element instanceof PySubscriptionExpression) {
final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression) element;
final PyExpression operand = subscriptionExpr.getOperand();
final PyExpression indexExpr = subscriptionExpr.getIndexExpression();
final PyType operandType = getType(operand, context);
if (operandType instanceof PyClassType) {
final PyClass cls = ((PyClassType) operandType).getPyClass();
final List<PyType> indexTypes = getIndexTypes(subscriptionExpr, context);
if (PyNames.TUPLE.equals(cls.getQualifiedName())) {
if (indexExpr instanceof PyTupleExpression) {
final PyExpression[] elements = ((PyTupleExpression) indexExpr).getElements();
if (elements.length == 2 && isEllipsis(elements[1])) {
return PyTupleType.createHomogeneous(element, indexTypes.get(0));
}
}
return PyTupleType.create(element, indexTypes);
} else if (indexExpr != null) {
return new PyCollectionTypeImpl(cls, false, indexTypes);
}
}
}
return null;
}
@Nullable
private static PyType getCallableType(@NotNull PsiElement resolved, @NotNull Context context) {
if (resolved instanceof PySubscriptionExpression) {
final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression) resolved;
final PyExpression operand = subscriptionExpr.getOperand();
final Collection<String> operandNames =
resolveToQualifiedNames(operand, context.getTypeContext());
if (operandNames.contains("typing.Callable")) {
final PyExpression indexExpr = subscriptionExpr.getIndexExpression();
if (indexExpr instanceof PyTupleExpression) {
final PyTupleExpression tupleExpr = (PyTupleExpression) indexExpr;
final PyExpression[] elements = tupleExpr.getElements();
if (elements.length == 2) {
final PyExpression parametersExpr = elements[0];
final PyExpression returnTypeExpr = elements[1];
if (parametersExpr instanceof PyListLiteralExpression) {
final List<PyCallableParameter> parameters = new ArrayList<>();
final PyListLiteralExpression listExpr = (PyListLiteralExpression) parametersExpr;
for (PyExpression argExpr : listExpr.getElements()) {
parameters.add(new PyCallableParameterImpl(null, getType(argExpr, context)));
}
final PyType returnType = getType(returnTypeExpr, context);
return new PyCallableTypeImpl(parameters, returnType);
}
if (isEllipsis(parametersExpr)) {
return new PyCallableTypeImpl(null, getType(returnTypeExpr, context));
}
}
}
}
}
return null;
}
@Nullable
private static PyType getType(@NotNull PyExpression expression, @NotNull Context context) {
final List<PyType> members = Lists.newArrayList();
for (PsiElement resolved : tryResolving(expression, context.getTypeContext())) {
members.add(getTypeForResolvedElement(resolved, context));
}
return PyUnionType.union(members);
}
@Nullable
private static PyType getGenericTypeBound(
@NotNull PyExpression[] typeVarArguments, @NotNull Context context) {
final List<PyType> types = new ArrayList<>();
for (int i = 1; i < typeVarArguments.length; i++) {
types.add(getType(typeVarArguments[i], context));
}
return PyUnionType.union(types);
}
public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
if (!TypeEvalStack.mayEvaluate(this)) {
return null;
}
try {
final boolean qualified = isQualified();
if (!qualified) {
String name = getReferencedName();
if (PyNames.NONE.equals(name)) {
return PyNoneType.INSTANCE;
}
}
PyType type = getTypeFromProviders(context);
if (type != null) {
return type;
}
if (qualified) {
PyType maybe_type = PyUtil.getSpecialAttributeType(this, context);
if (maybe_type != null) return maybe_type;
Ref<PyType> typeOfProperty = getTypeOfProperty(context);
if (typeOfProperty != null) {
return typeOfProperty.get();
}
}
final PsiPolyVariantReference reference =
getReference(PyResolveContext.noImplicits().withTypeEvalContext(context));
final List<PsiElement> targets = PyUtil.multiResolveTopPriority(reference);
if (targets.isEmpty()) {
return getQualifiedReferenceTypeByControlFlow(context);
}
final List<PyType> members = new ArrayList<PyType>();
for (PsiElement target : targets) {
if (target == this || target == null) {
continue;
}
if (!target.isValid()) {
LOG.error(
"Reference "
+ this
+ " resolved to invalid element "
+ target
+ " (text="
+ target.getText()
+ ")");
continue;
}
members.add(getTypeFromTarget(target, context, this));
}
return PyUnionType.union(members);
} finally {
TypeEvalStack.evaluated(this);
}
}
@NotNull
public QualifiedResolveResult followAssignmentsChain(PyResolveContext resolveContext) {
PyReferenceExpression seeker = this;
QualifiedResolveResult ret = null;
List<PyExpression> qualifiers = new ArrayList<PyExpression>();
PyExpression qualifier = seeker.getQualifier();
if (qualifier != null) {
qualifiers.add(qualifier);
}
Set<PsiElement> visited = new HashSet<PsiElement>();
visited.add(this);
SEARCH:
while (ret == null) {
ResolveResult[] targets = seeker.getReference(resolveContext).multiResolve(false);
for (ResolveResult target : targets) {
PsiElement elt = target.getElement();
if (elt instanceof PyTargetExpression) {
PsiElement assigned_from = null;
final PyTargetExpression expr = (PyTargetExpression) elt;
final TypeEvalContext context = resolveContext.getTypeEvalContext();
if (context.maySwitchToAST(expr) || expr.getStub() == null) {
assigned_from = expr.findAssignedValue();
}
// TODO: Maybe findAssignedValueByStub() should become a part of the PyTargetExpression
// interface
else if (elt instanceof PyTargetExpressionImpl) {
assigned_from = ((PyTargetExpressionImpl) elt).findAssignedValueByStub(context);
}
if (assigned_from instanceof PyReferenceExpression) {
if (visited.contains(assigned_from)) {
break;
}
visited.add(assigned_from);
seeker = (PyReferenceExpression) assigned_from;
if (seeker.getQualifier() != null) {
qualifiers.add(seeker.getQualifier());
}
continue SEARCH;
} else if (assigned_from != null)
ret = new QualifiedResolveResultImpl(assigned_from, qualifiers, false);
} else if (ret == null && elt instanceof PyElement && target.isValidResult()) {
// remember this result, but a further reference may be the next resolve result
ret =
new QualifiedResolveResultImpl(
elt, qualifiers, target instanceof ImplicitResolveResult);
}
}
// all resolve results checked, reassignment not detected, nothing more to do
break;
}
if (ret == null) ret = EMPTY_RESULT;
return ret;
}
@NotNull
private static List<PyType> getIndexTypes(
@NotNull PySubscriptionExpression expression, @NotNull Context context) {
final List<PyType> types = new ArrayList<>();
final PyExpression indexExpr = expression.getIndexExpression();
if (indexExpr instanceof PyTupleExpression) {
final PyTupleExpression tupleExpr = (PyTupleExpression) indexExpr;
for (PyExpression expr : tupleExpr.getElements()) {
types.add(getType(expr, context));
}
} else if (indexExpr != null) {
types.add(getType(indexExpr, context));
}
return types;
}
@Nullable
private static PyType getGenericConstructorType(
@NotNull PyFunction function, @NotNull Context context) {
if (PyUtil.isInit(function)) {
final PyClass cls = function.getContainingClass();
if (cls != null) {
final List<PyGenericType> genericTypes = collectGenericTypes(cls, context);
final List<PyType> elementTypes = new ArrayList<>(genericTypes);
if (!elementTypes.isEmpty()) {
return new PyCollectionTypeImpl(cls, false, elementTypes);
}
}
}
return null;
}
@Nullable
public Ref<PyType> getParameterType(
@NotNull PyNamedParameter param, @NotNull PyFunction func, @NotNull TypeEvalContext context) {
final PyAnnotation annotation = param.getAnnotation();
if (annotation != null) {
// XXX: Requires switching from stub to AST
final PyExpression value = annotation.getValue();
if (value != null) {
final PyType type = getType(value, new Context(context));
if (type != null) {
final PyType optionalType = getOptionalTypeFromDefaultNone(param, type, context);
return Ref.create(optionalType != null ? optionalType : type);
}
}
}
final String paramComment = param.getTypeCommentAnnotation();
if (paramComment != null) {
return Ref.create(getStringBasedType(paramComment, param, new Context(context)));
}
final String comment = func.getTypeCommentAnnotation();
if (comment != null) {
final PyTypeParser.ParseResult result =
PyTypeParser.parsePep484FunctionTypeComment(param, comment);
final PyCallableType functionType = as(result.getType(), PyCallableType.class);
if (functionType != null) {
final List<PyCallableParameter> paramTypes = functionType.getParameters(context);
// Function annotation of kind (...) -> Type
if (paramTypes == null) {
return Ref.create();
}
final PyParameter[] funcParams = func.getParameterList().getParameters();
final int startOffset = omitFirstParamInTypeComment(func) ? 1 : 0;
for (int paramIndex = 0; paramIndex < funcParams.length; paramIndex++) {
if (funcParams[paramIndex] == param) {
final int typeIndex = paramIndex - startOffset;
if (typeIndex >= 0 && typeIndex < paramTypes.size()) {
return Ref.create(paramTypes.get(typeIndex).getType(context));
}
break;
}
}
}
}
return null;
}
/** @param self should be this */
@NotNull
private static List<PyAssignmentStatement> findAttributesStatic(@NotNull final PsiElement self) {
final List<PyAssignmentStatement> result = new ArrayList<>();
for (final PyAssignmentStatement statement :
new PsiQuery(self).siblings(PyAssignmentStatement.class).getElements()) {
for (final PyQualifiedExpression targetExpression :
new PsiQuery(statement.getTargets()).filter(PyQualifiedExpression.class).getElements()) {
final PyExpression qualifier = targetExpression.getQualifier();
if (qualifier == null) {
continue;
}
final PsiReference qualifierReference = qualifier.getReference();
if (qualifierReference == null) {
continue;
}
if (qualifierReference.isReferenceTo(self)) {
result.add(statement);
}
}
}
return result;
}
private static PyType getTypeByControlFlow(
@NotNull String name,
@NotNull TypeEvalContext context,
@NotNull PyExpression anchor,
@NotNull ScopeOwner scopeOwner) {
PyAugAssignmentStatement augAssignment =
PsiTreeUtil.getParentOfType(anchor, PyAugAssignmentStatement.class);
try {
final PyElement element = augAssignment != null ? augAssignment : anchor;
final List<ReadWriteInstruction> defs =
PyDefUseUtil.getLatestDefs(scopeOwner, name, element, true);
if (!defs.isEmpty()) {
PyType type = defs.get(0).getType(context, anchor);
for (int i = 1; i < defs.size(); i++) {
type = PyUnionType.union(type, defs.get(i).getType(context, anchor));
}
return type;
}
} catch (PyDefUseUtil.InstructionNotFoundException ignored) {
}
return null;
}
@Nullable
private Modifier getWrappersFromStub() {
final StubElement parentStub = getStub().getParentStub();
final List childrenStubs = parentStub.getChildrenStubs();
int index = childrenStubs.indexOf(getStub());
if (index >= 0 && index < childrenStubs.size() - 1) {
StubElement nextStub = (StubElement) childrenStubs.get(index + 1);
if (nextStub instanceof PyTargetExpressionStub) {
final PyTargetExpressionStub targetExpressionStub = (PyTargetExpressionStub) nextStub;
if (targetExpressionStub.getInitializerType()
== PyTargetExpressionStub.InitializerType.CallExpression) {
final QualifiedName qualifiedName = targetExpressionStub.getInitializer();
if (QualifiedName.fromComponents(PyNames.CLASSMETHOD).equals(qualifiedName)) {
return CLASSMETHOD;
}
if (QualifiedName.fromComponents(PyNames.STATICMETHOD).equals(qualifiedName)) {
return STATICMETHOD;
}
}
}
}
return null;
}
@NotNull
private static List<PsiElement> tryResolving(
@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
final List<PsiElement> elements = Lists.newArrayList();
if (expression instanceof PyReferenceExpression) {
final PyReferenceExpression referenceExpr = (PyReferenceExpression) expression;
final PyResolveContext resolveContext =
PyResolveContext.noImplicits().withTypeEvalContext(context);
final PsiPolyVariantReference reference = referenceExpr.getReference(resolveContext);
final List<PsiElement> resolved = PyUtil.multiResolveTopPriority(reference);
for (PsiElement element : resolved) {
if (element instanceof PyFunction) {
final PyFunction function = (PyFunction) element;
if (PyUtil.isInit(function)) {
final PyClass cls = function.getContainingClass();
if (cls != null) {
elements.add(cls);
continue;
}
}
} else if (element instanceof PyTargetExpression) {
final PyTargetExpression targetExpr = (PyTargetExpression) element;
// XXX: Requires switching from stub to AST
final PyExpression assignedValue = targetExpr.findAssignedValue();
if (assignedValue != null) {
elements.add(assignedValue);
continue;
}
}
if (element != null) {
elements.add(element);
}
}
}
return !elements.isEmpty() ? elements : Collections.singletonList(expression);
}
@Nullable
private String extractReturnType() {
final String ARROW = "->";
final StructuredDocString structuredDocString = getStructuredDocString();
if (structuredDocString != null) {
return structuredDocString.getReturnType();
}
final String docString = getDocStringValue();
if (docString != null && docString.contains(ARROW)) {
final List<String> lines = StringUtil.split(docString, "\n");
while (lines.size() > 0 && lines.get(0).trim().length() == 0) {
lines.remove(0);
}
if (lines.size() > 1 && lines.get(1).trim().length() == 0) {
String firstLine = lines.get(0);
int pos = firstLine.lastIndexOf(ARROW);
if (pos >= 0) {
return firstLine.substring(pos + 2).trim();
}
}
}
return null;
}