@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;
}
@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;
}
@Nullable
private static PyType getTypeFromTargetExpression(
@NotNull PyTargetExpression expression, @NotNull Context context) {
// XXX: Requires switching from stub to AST
final PyExpression assignedValue = expression.findAssignedValue();
return assignedValue != null ? getTypeForResolvedElement(assignedValue, context) : 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);
}
@Override
protected PsiElement checkLocalScope() {
return myTarget.getContainingFile();
}