Fixed non-determinism for target exprs with qualified assigned values during stubs->AST (PY-15558)

Now we use exactly the same algorithm for resolving values assigned to
target expressions for both stubs and AST.

python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java

@@ -15,11 +15,9 @@
  */
 package com.jetbrains.python.psi;
 
-import com.intellij.psi.PsiNameIdentifierOwner;
-import com.intellij.psi.PsiNamedElement;
-import com.intellij.psi.PsiReference;
-import com.intellij.psi.StubBasedPsiElement;
+import com.intellij.psi.*;
 import com.intellij.psi.util.QualifiedName;
+import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -36,11 +34,26 @@ public interface PyTargetExpression extends PyQualifiedExpression, PsiNamedEleme
    * Find the value that maps to this target expression in an enclosing assignment expression.
    * Does not work with other expressions (e.g. if the target is in a 'for' loop).
    *
+   * Operates at the AST level.
+   *
    * @return the expression assigned to target via an enclosing assignment expression, or null.
    */
   @Nullable
   PyExpression findAssignedValue();
 
+  /**
+   * Resolves the value that maps to this target expression in an enclosing assignment expression.
+   *
+   * This method does not access AST if underlying PSI is stub based and the context doesn't allow switching to AST.
+   */
+  @Nullable
+  PsiElement resolveAssignedValue(@NotNull PyResolveContext resolveContext);
+
+  /**
+   * Returns the qualified name (if there is any) assigned to the expression.
+   *
+   * This method does not access AST if underlying PSI is stub based.
+   */
   @Nullable
   QualifiedName getAssignedQName();
 

python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java

@@ -144,15 +144,14 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
       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) {
+          final PsiElement assigned_from;
+          if (context.maySwitchToAST(expr)) {
             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);
+          else {
+            assigned_from = expr.resolveAssignedValue(resolveContext);
           }
           if (assigned_from instanceof PyReferenceExpression) {
             if (visited.contains(assigned_from)) {

python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java

@@ -144,17 +144,13 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
         return type;
       }
       if (!context.maySwitchToAST(this)) {
-        final PsiElement value = getStub() != null ? findAssignedValueByStub(context) : findAssignedValue();
+        final PsiElement value = resolveAssignedValue(PyResolveContext.noImplicits().withTypeEvalContext(context));
         if (value instanceof PyTypedElement) {
           type = context.getType((PyTypedElement)value);
           if (type instanceof PyNoneType) {
             return null;
           }
-          if (type instanceof PyFunctionTypeImpl) {
-            return type;
-          }
-          // We are unsure about the type since it may be inferred from the stub based on incomplete information
-          return PyUnionType.createWeakType(type);
+          return type;
         }
         return null;
       }
@@ -476,6 +472,50 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
   }
 
   @Nullable
+  @Override
+  public PsiElement resolveAssignedValue(@NotNull PyResolveContext resolveContext) {
+    final TypeEvalContext context = resolveContext.getTypeEvalContext();
+    if (context.maySwitchToAST(this)) {
+      final PyExpression value = findAssignedValue();
+      if (value != null) {
+        final List<PsiElement> results = PyUtil.multiResolveTopPriority(value, resolveContext);
+        return !results.isEmpty() ? results.get(0) : null;
+      }
+      return null;
+    }
+    else {
+      final QualifiedName qName = getAssignedQName();
+      if (qName != null) {
+        final ScopeOwner owner = ScopeUtil.getScopeOwner(this);
+        if (owner instanceof PyTypedElement) {
+          final List<String> components = qName.getComponents();
+          if (!components.isEmpty()) {
+            PsiElement resolved = owner;
+            for (String component : components) {
+              if (!(resolved instanceof PyTypedElement)) {
+                return null;
+              }
+              final PyType qualifierType = context.getType((PyTypedElement)resolved);
+              if (qualifierType == null) {
+                return null;
+              }
+              final List<? extends RatedResolveResult> results = qualifierType.resolveMember(component, null, AccessDirection.READ,
+                                                                                             resolveContext);
+              if (results == null || results.isEmpty()) {
+                return null;
+              }
+              resolved = results.get(0).getElement();
+            }
+            return resolved;
+          }
+        }
+      }
+      return null;
+    }
+  }
+
+  @Nullable
+  @Override
   public PyExpression findAssignedValue() {
     if (isValid()) {
       PyAssignmentStatement assignment = PsiTreeUtil.getParentOfType(this, PyAssignmentStatement.class);
@@ -490,6 +530,8 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
     return null;
   }
 
+  @Nullable
+  @Override
   public QualifiedName getAssignedQName() {
     final PyTargetExpressionStub stub = getStub();
     if (stub != null) {
@@ -501,35 +543,6 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
     return PyPsiUtils.asQualifiedName(findAssignedValue());
   }
 
-  @Nullable
-  public PsiElement findAssignedValueByStub(@NotNull TypeEvalContext context) {
-    final PyTargetExpressionStub stub = getStub();
-    if (stub != null && stub.getInitializerType() == PyTargetExpressionStub.InitializerType.ReferenceExpression) {
-      final QualifiedName initializer = stub.getInitializer();
-      // TODO: Support qualified stub initializers
-      if (initializer != null && initializer.getComponentCount() == 1) {
-        final String name = initializer.getLastComponent();
-        if (name != null) {
-          final PsiElement parent = getParentByStub();
-          if (parent instanceof PyFile) {
-            return ((PyFile)parent).getElementNamed(name);
-          }
-          else if (parent instanceof PyClass) {
-            final PyType type = context.getType((PyClass)parent);
-            if (type != null) {
-              final List<? extends RatedResolveResult> results = type.resolveMember(name, null, AccessDirection.READ,
-                                                                                    PyResolveContext.noImplicits());
-              if (results != null && !results.isEmpty()) {
-                return results.get(0).getElement();
-              }
-            }
-          }
-        }
-      }
-    }
-    return null;
-  }
-
   @Override
   public QualifiedName getCalleeName() {
     final PyTargetExpressionStub stub = getStub();