PY-45206 Resolve instance attributes found in the same scope using control flow

We have to define a resolve order so that the attribute is never defined in terms of itself. The main result is preventing recursion in the type checker in case of an infinite loop. It used to cause unstable type inference based on the order in which we process elements.

A side effect of this change is a more precise behaviour of resolve for attributes where they are clearly not defined yet like:

    class C:
        def f(self):
            print(self.foo)  # foo is now unresolved
            self.foo = 0

See more examples in the unit tests for this fix.

GitOrigin-RevId: a791ed9a16df64935bf70e2933428c0950e5e7d3

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyClassImpl.java

@@ -1245,19 +1245,17 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
 
   @Override
   public boolean processInstanceLevelDeclarations(@NotNull PsiScopeProcessor processor, @Nullable PsiElement location) {
-    final Map<String, PyTargetExpression> declarationsInMethod = new HashMap<>();
     final PyFunction instanceMethod = PsiTreeUtil.getStubOrPsiParentOfType(location, PyFunction.class);
     final PyClass containingClass = instanceMethod != null ? instanceMethod.getContainingClass() : null;
     if (instanceMethod != null && containingClass != null && CompletionUtilCoreImpl.getOriginalElement(containingClass) == this) {
-      collectInstanceAttributes(instanceMethod, declarationsInMethod);
-      for (PyTargetExpression targetExpression : declarationsInMethod.values()) {
-        if (!processor.execute(targetExpression, ResolveState.initial())) {
+      for (PyTargetExpression target : getTargetExpressions(instanceMethod)) {
+        if (PyUtil.isInstanceAttribute(target) && !processor.execute(target, ResolveState.initial())) {
           return false;
         }
       }
     }
     for (PyTargetExpression expr : getInstanceAttributes()) {
-      if (declarationsInMethod.containsKey(expr.getName())) {
+      if (instanceMethod != null && ScopeUtil.getScopeOwner(expr) == instanceMethod) {
         continue;
       }
       if (!processor.execute(expr, ResolveState.initial())) return false;

python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java

@@ -137,7 +137,7 @@ public class PyQualifiedReference extends PyReferenceImpl {
       return false;
     }
     for (PyExpression ex : collectAssignedAttributes(qName, qualifier)) {
-      if (referencedName.equals(ex.getName())) {
+      if (referencedName.equals(ex.getName()) && !PyUtil.isInstanceAttribute(ex)) {
         ret.poke(ex, RatedResolveResult.RATE_NORMAL);
         return true;
       }

python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveProcessor.java

@@ -29,7 +29,7 @@ public class PyResolveProcessor implements PsiScopeProcessor {
   private final boolean myLocalResolve;
   @NotNull private final Map<PsiElement, PyImportedNameDefiner> myResults = Maps.newLinkedHashMap();
   @NotNull private final Map<PsiElement, PyImportedNameDefiner> myImplicitlyImportedResults = Maps.newLinkedHashMap();
-  @Nullable private ScopeOwner myOwner;
+  @Nullable protected ScopeOwner myOwner;
 
   public PyResolveProcessor(@NotNull String name) {
     this(name, false);
@@ -99,21 +99,25 @@ public class PyResolveProcessor implements PsiScopeProcessor {
     return definer.multiResolveName(myName);
   }
 
-  private boolean tryAddResult(@Nullable PsiElement element, @Nullable PyImportedNameDefiner definer) {
+  protected boolean tryAddResult(@Nullable PsiElement element, @Nullable PyImportedNameDefiner definer) {
     final ScopeOwner owner = ScopeUtil.getScopeOwner(definer != null ? definer : element);
     if (myOwner == null) {
       myOwner = owner;
     }
     final boolean sameScope = owner == myOwner;
     if (sameScope) {
-      // XXX: In 'from foo import foo' inside __init__.py the preferred result is explicitly imported 'foo'
-      if (definer instanceof PyFromImportStatement) {
-        myImplicitlyImportedResults.put(element, definer);
-      }
-      else {
-        myResults.put(element, definer);
-      }
+      addResult(element, definer);
     }
     return sameScope;
   }
+
+  protected final void addResult(@Nullable PsiElement element, @Nullable PyImportedNameDefiner definer) {
+    // XXX: In 'from foo import foo' inside __init__.py the preferred result is explicitly imported 'foo'
+    if (definer instanceof PyFromImportStatement) {
+      myImplicitlyImportedResults.put(element, definer);
+    }
+    else {
+      myResults.put(element, definer);
+    }
+  }
 }

python/python-psi-impl/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java

@@ -8,12 +8,15 @@ import com.intellij.psi.*;
 import com.intellij.psi.scope.PsiScopeProcessor;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ArrayUtilRt;
+import com.intellij.util.ObjectUtils;
 import com.intellij.util.ProcessingContext;
 import com.intellij.util.Processor;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.PyCustomMember;
 import com.jetbrains.python.codeInsight.PyCustomMemberUtils;
+import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyCallExpressionHelper;
@@ -22,6 +25,7 @@ import com.jetbrains.python.psi.impl.ResolveResultList;
 import com.jetbrains.python.psi.impl.references.PyReferenceImpl;
 import com.jetbrains.python.psi.resolve.*;
 import com.jetbrains.python.pyi.PyiUtil;
+import com.jetbrains.python.refactoring.PyDefUseUtil;
 import com.jetbrains.python.toolbox.Maybe;
 import one.util.streamex.EntryStream;
 import one.util.streamex.StreamEx;
@@ -445,7 +449,7 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
                                                                  @NotNull String name,
                                                                  @Nullable PyExpression location,
                                                                  @NotNull TypeEvalContext context) {
-    final PyResolveProcessor processor = new PyResolveProcessor(name);
+    final PyAttributesProcessor processor = new PyAttributesProcessor(name, location);
     final Map<PsiElement, PyImportedNameDefiner> results;
 
     if (isDefinition || cls.processInstanceLevelDeclarations(processor, location)) {
@@ -754,4 +758,74 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
     }
     return new PyClassTypeImpl(pyClass, isDefinition);
   }
+
+  /**
+   * <p>Control flow aware Python attributes resolver.</p>
+   *
+   * <p>It respects control flow if a resolve candidate is defined in the same scope as the location we resolve the attribute from.</p>
+   *
+   * <p>Since an attribute doesn't have to be defined in the same method we use it, we have to assume that an attribute we cannot
+   * resolve via the control flow graph is defined in some other method. If the attribute is not resolved via the graph, but is defined
+   * in a sibling if-elif-else branch, we assume it will become available in our branch eventually in subsequent method calls.</p>
+   */
+  private static final class PyAttributesProcessor extends PyResolveProcessor {
+    @Nullable private final PyExpression myLocation;
+
+    PyAttributesProcessor(@NotNull String name, @Nullable PyExpression location) {
+      super(name);
+      myLocation = location;
+    }
+
+    @Override
+    protected boolean tryAddResult(@Nullable PsiElement element, @Nullable PyImportedNameDefiner definer) {
+      PsiElement psiElement = definer != null ? definer : element;
+      if (inSameScope(psiElement, myLocation)) {
+        if (PsiTreeUtil.isAncestor(psiElement, myLocation, false) ||
+            PyDefUseUtil.isDefinedBefore(psiElement, myLocation) ||
+            inDifferentBranchesOfSameIfStatement(psiElement, myLocation)) {
+          if (myOwner == null) {
+            myOwner = ScopeUtil.getScopeOwner(psiElement);
+          }
+          addResult(element, definer);
+        }
+        return true;
+      }
+      return super.tryAddResult(element, definer);
+    }
+
+    private static boolean inSameScope(@Nullable PsiElement e1, @Nullable PsiElement e2) {
+      if (e1 == null || e2 == null) return false;
+      ScopeOwner o1 = ScopeUtil.getScopeOwner(e1);
+      ScopeOwner o2 = ScopeUtil.getScopeOwner(e2);
+      return o1 != null && o1 == o2;
+    }
+
+    private static boolean inDifferentBranchesOfSameIfStatement(@NotNull PsiElement e1, @NotNull PsiElement e2) {
+      PyIfStatement ifStatement = ObjectUtils.tryCast(PsiTreeUtil.findCommonParent(e1, e2), PyIfStatement.class);
+      if (ifStatement == null) return false;
+      List<PyStatementPart> parts = getIfStatementParts(ifStatement);
+      PyStatementPart p1 = findIfStatementPartByElement(e1, parts);
+      PyStatementPart p2 = findIfStatementPartByElement(e2, parts);
+      return p1 != p2;
+    }
+
+    private static PyStatementPart findIfStatementPartByElement(@NotNull PsiElement element, @NotNull List<PyStatementPart> parts) {
+      return StreamEx.of(parts)
+        .filter(part -> PsiTreeUtil.isAncestor(part, element, true))
+        .findFirst()
+        .orElse(null);
+    }
+
+    @NotNull
+    private static List<PyStatementPart> getIfStatementParts(@NotNull PyIfStatement statement) {
+      List<PyStatementPart> parts = new ArrayList<>();
+      parts.add(statement.getIfPart());
+      parts.addAll(Arrays.asList(statement.getElifParts()));
+      PyElsePart elsePart = statement.getElsePart();
+      if (elsePart != null) {
+        parts.add(elsePart);
+      }
+      return parts;
+    }
+  }
 }

python/testData/resolve/InstanceAttrAbove.py

@@ -0,0 +1,5 @@
+class C:
+    def f(self):
+        self.foo = 1
+        return self.foo
+#                   <ref>

python/testData/resolve/InstanceAttrBelowEarlierByControlFlow.py

@@ -0,0 +1,10 @@
+class C:
+    def f(self):
+        c = False
+        while True:
+            if c:
+                return self.foo
+                #           <ref>
+            else:
+                c = True
+                self.foo = 1

python/testData/resolve/InstanceAttrBothEarlierAndLater.py

@@ -0,0 +1,6 @@
+class C:
+    def f(self):
+        self.foo = 1
+        if self.foo:
+            #   <ref>
+            self.foo = 0

python/testData/resolve/InstanceAttrInheritedAndAbove.py

@@ -0,0 +1,10 @@
+class B:
+    def g(self):
+        self.foo = 0
+
+
+class C(B):
+    def f(self):
+        self.foo = 1
+        return self.foo
+        #           <ref>

python/testData/resolve/InstanceAttrInheritedAndBelow.py

@@ -0,0 +1,11 @@
+class B:
+    def g(self):
+        self.foo = 0
+
+
+class C(B):
+    def f(self):
+        x = self.foo
+        #        <ref>
+        self.foo = 1
+        return x

python/testData/resolve/InstanceAttrOtherMethod.py

@@ -0,0 +1,7 @@
+class C:
+    def f(self):
+        return self.foo
+    #               <ref>
+
+    def g(self):
+        self.foo = 1

python/testData/resolve/InstanceAttrOtherMethodAndAbove.py

@@ -0,0 +1,8 @@
+class C:
+    def g(self):
+        self.foo = 0
+
+    def f(self):
+        self.foo = 1
+        return self.foo
+        #           <ref>

python/testData/resolve/InstanceAttrOtherMethodAndBelow.py

@@ -0,0 +1,9 @@
+class C:
+    def f(self):
+        x = self.foo
+        #        <ref>
+        self.foo = 1
+        return x
+
+    def g(self):
+        self.foo = 0

python/testData/resolve/NoResolveInstanceAttrBelow.py

@@ -0,0 +1,5 @@
+class C:
+    def f(self):
+        x = self.foo
+        #        <ref>
+        self.foo = 1

python/testData/resolve/NoResolveInstanceAttrSameLine.py

@@ -0,0 +1,4 @@
+class C:
+    def f(self):
+        self.foo = [1, 2, self.foo]
+        #                      <ref>

python/testSrc/com/jetbrains/python/Py3ResolveTest.java

@@ -17,6 +17,8 @@ package com.jetbrains.python;
 
 import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.ObjectUtils;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.fixtures.PyResolveTestCase;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
@@ -775,4 +777,57 @@ public class Py3ResolveTest extends PyResolveTestCase {
   public void testTypeVarClassObjectBoundAttribute() {
     assertNull(doResolve());
   }
+
+  public void testInstanceAttrAbove() {
+    assertResolvesTo(PyTargetExpression.class, "foo");
+  }
+
+  public void testNoResolveInstanceAttrBelow() {
+    assertUnresolved();
+  }
+
+  public void testNoResolveInstanceAttrSameLine() {
+    assertUnresolved();
+  }
+
+  public void testInstanceAttrOtherMethod() {
+    assertResolvesTo(PyTargetExpression.class, "foo");
+  }
+
+  public void testInstanceAttrOtherMethodAndAbove() {
+    final PyTargetExpression target = assertResolvesTo(PyTargetExpression.class, "foo");
+    final PyFunction function = ObjectUtils.tryCast(ScopeUtil.getScopeOwner(target), PyFunction.class);
+    assertNotNull(function);
+    assertEquals("f", function.getName());
+  }
+
+  public void testInstanceAttrOtherMethodAndBelow() {
+    final PyTargetExpression target = assertResolvesTo(PyTargetExpression.class, "foo");
+    final PyFunction function = ObjectUtils.tryCast(ScopeUtil.getScopeOwner(target), PyFunction.class);
+    assertNotNull(function);
+    assertEquals("g", function.getName());
+  }
+
+  public void testInstanceAttrInheritedAndAbove() {
+    final PyTargetExpression target = assertResolvesTo(PyTargetExpression.class, "foo");
+    final PyFunction function = ObjectUtils.tryCast(ScopeUtil.getScopeOwner(target), PyFunction.class);
+    assertNotNull(function);
+    assertEquals("f", function.getName());
+  }
+
+  public void testInstanceAttrInheritedAndBelow() {
+    final PyTargetExpression target = assertResolvesTo(PyTargetExpression.class, "foo");
+    final PyFunction function = ObjectUtils.tryCast(ScopeUtil.getScopeOwner(target), PyFunction.class);
+    assertNotNull(function);
+    assertEquals("g", function.getName());
+  }
+
+  public void testInstanceAttrBelowEarlierByControlFlow() {
+    assertResolvesTo(PyTargetExpression.class, "foo");
+  }
+
+  public void testInstanceAttrBothEarlierAndLater() {
+    PyTargetExpression target = assertResolvesTo(PyTargetExpression.class, "foo");
+    assertEquals("self.foo = 1", target.getParent().getText());
+  }
 }

python/testSrc/com/jetbrains/python/Py3TypeTest.java

@@ -4,6 +4,7 @@ package com.jetbrains.python;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiFile;
 import com.jetbrains.python.fixtures.PyTestCase;
+import com.jetbrains.python.inspections.PyTypeCheckerInspectionTest;
 import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
@@ -1071,9 +1072,43 @@ public class Py3TypeTest extends PyTestCase {
            "        expr = m");
   }
 
+  /**
+   * @see #testRecursiveDictTopDown()
+   * @see PyTypeCheckerInspectionTest#testRecursiveDictAttribute()
+   */
+  public void testRecursiveDictBottomUp() {
+    String text = "class C:\n" +
+                  "    def f(self, x):\n" +
+                  "        self.foo = x\n" +
+                  "        self.foo = {'foo': self.foo}\n" +
+                  "        expr = self.foo\n";
+    myFixture.configureByText(PythonFileType.INSTANCE, text);
+    PyExpression dict = myFixture.findElementByText("{'foo': self.foo}", PyExpression.class);
+    assertExpressionType("Dict[str, Any]", dict);
+    final PyExpression expr = myFixture.findElementByText("expr", PyExpression.class);
+    assertExpressionType("Dict[str, Any]", expr);
+  }
+
+  public void testRecursiveDictTopDown() {
+    String text = "class C:\n" +
+                  "    def f(self, x):\n" +
+                  "        self.foo = x\n" +
+                  "        self.foo = {'foo': self.foo}\n" +
+                  "        expr = self.foo\n";
+    myFixture.configureByText(PythonFileType.INSTANCE, text);
+    final PyExpression expr = myFixture.findElementByText("expr", PyExpression.class);
+    assertExpressionType("Dict[str, Any]", expr);
+    PyExpression dict = myFixture.findElementByText("{'foo': self.foo}", PyExpression.class);
+    assertExpressionType("Dict[str, Any]", dict);
+  }
+
   private void doTest(final String expectedType, final String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);
     final PyExpression expr = myFixture.findElementByText("expr", PyExpression.class);
+    assertExpressionType(expectedType, expr);
+  }
+
+  private void assertExpressionType(String expectedType, PyExpression expr) {
     final Project project = expr.getProject();
     final PsiFile containingFile = expr.getContainingFile();
     assertType(expectedType, expr, TypeEvalContext.codeAnalysis(project, containingFile));

python/testSrc/com/jetbrains/python/PyTypeTest.java

@@ -2487,7 +2487,7 @@ public class PyTypeTest extends PyTestCase {
 
   // PY-21175
   public void testLazyAttributeInitialization() {
-    doTest("int",
+    doTest("Union[int, Any]",
            "class C:\n" +
            "    def __init__(self):\n" +
            "        self.attr = None\n" +