Fixed resolve and completion for locally assigned instance fields (PY-4279)

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

@@ -1,6 +1,8 @@
 package com.jetbrains.python.psi.impl;
 
+import com.intellij.codeInsight.completion.CompletionUtil;
 import com.intellij.lang.ASTNode;
+import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.NotNullLazyValue;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
@@ -832,12 +834,12 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     // __init__ takes priority over all other methods
     PyFunctionImpl initMethod = (PyFunctionImpl)findMethodByName(PyNames.INIT, false);
     if (initMethod != null) {
-      collectInstanceAttributes(initMethod, result);
+      collectInstanceAttributes(initMethod, result, null);
     }
     final PyFunction[] methods = getMethods();
     for (PyFunction method : methods) {
       if (!PyNames.INIT.equals(method.getName())) {
-        collectInstanceAttributes(method, result);
+        collectInstanceAttributes(method, result, null);
       }
     }
 
@@ -845,7 +847,9 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     return new ArrayList<PyTargetExpression>(expressions);
   }
 
-  private static void collectInstanceAttributes(PyFunction method, final Map<String, PyTargetExpression> result) {
+  private static void collectInstanceAttributes(@NotNull PyFunction method,
+                                                @NotNull final Map<String, PyTargetExpression> result,
+                                                @Nullable PsiElement anchor) {
     final PyParameter[] params = method.getParameterList().getParameters();
     if (params.length == 0) {
       return;
@@ -859,25 +863,49 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
         }
       }
     }
+    else if (anchor != null) {
+      PyResolveUtil.treeCrawlUp(new PsiScopeProcessor() {
+        @Override
+        public boolean execute(PsiElement element, ResolveState state) {
+          if (element instanceof PyAssignmentStatement) {
+            collectNewTargets(result, (PyAssignmentStatement)element);
+          }
+          return true;
+        }
+
+        @Nullable
+        @Override
+        public <T> T getHint(Key<T> hintKey) {
+          return null;
+        }
+
+        @Override
+        public void handleEvent(Event event, @Nullable Object associated) {
+        }
+      }, false, anchor, method);
+    }
     else {
-      // NOTE: maybe treeCrawlUp would be more precise, but currently it works well enough; don't care.
       final PyStatementList statementList = method.getStatementList();
       if (statementList != null) {
         statementList.accept(new PyRecursiveElementVisitor() {
           public void visitPyAssignmentStatement(final PyAssignmentStatement node) {
             super.visitPyAssignmentStatement(node);
-            final PyExpression[] targets = node.getTargets();
-            for (PyExpression target : targets) {
-              if (PyUtil.isInstanceAttribute(target) && !result.containsKey(target.getName())) {
-                result.put(target.getName(), (PyTargetExpression) target);
-              }
-            }
+            collectNewTargets(result, node);
           }
         });
       }
     }
   }
 
+  private static void collectNewTargets(Map<String, PyTargetExpression> collected, PyAssignmentStatement node) {
+    final PyExpression[] targets = node.getTargets();
+    for (PyExpression target : targets) {
+      if (target instanceof PyTargetExpression && PyUtil.isInstanceAttribute(target) && !collected.containsKey(target.getName())) {
+        collected.put(target.getName(), (PyTargetExpression)target);
+      }
+    }
+  }
+
   public boolean isNewStyleClass() {
     return myNewStyle.getValue().getValue();
   }
@@ -947,8 +975,9 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
   public boolean processInstanceLevelDeclarations(PsiScopeProcessor processor, @Nullable PyExpression location) {
     Map<String, PyTargetExpression> declarationsInMethod = new HashMap<String, PyTargetExpression>();
     PyFunction instanceMethod = PsiTreeUtil.getParentOfType(location, PyFunction.class);
-    if (instanceMethod != null && instanceMethod.getContainingClass() == this) {
-      collectInstanceAttributes(instanceMethod, declarationsInMethod);
+    final PyClass containingClass = instanceMethod != null ? instanceMethod.getContainingClass() : null;
+    if (instanceMethod != null && containingClass != null && CompletionUtil.getOriginalElement(containingClass) == this) {
+      collectInstanceAttributes(instanceMethod, declarationsInMethod, location);
       for (PyTargetExpression targetExpression : declarationsInMethod.values()) {
         if (!processor.execute(targetExpression, ResolveState.initial())) {
           return false;

python/testData/completion/fieldReassignment.after.py

@@ -0,0 +1,9 @@
+class C1(object):
+    def method1(self):
+        pass
+
+class Test(object):
+    def __init__(self, x):
+        self.x = x
+        self.x = C1()
+        self.x.method1()

python/testData/completion/fieldReassignment.py

@@ -0,0 +1,9 @@
+class C1(object):
+    def method1(self):
+        pass
+
+class Test(object):
+    def __init__(self, x):
+        self.x = x
+        self.x = C1()
+        self.x.meth<caret>

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

@@ -274,7 +274,25 @@ public class PyTypeTest extends PyTestCase {
            "        self.assertIsInstance(x, int)\n" +
            "        expr = x\n");
   }
-  
+
+  // PY-4279
+  public void testFieldReassignment() {
+    doTest("C1",
+           "class C1(object):\n" +
+           "    def m1(self):\n" +
+           "        pass\n" +
+           "\n" +
+           "class C2(object):\n" +
+           "    def m2(self):\n" +
+           "        pass\n" +
+           "\n" +
+           "class Test(object):\n" +
+           "    def __init__(self, param1):\n" +
+           "        self.x = param1\n" +
+           "        self.x = C1()\n" +
+           "        expr = self.x\n");
+  }
+
   public void testSOEOnRecursiveCall() {
     PyExpression expr = parseExpr("def foo(x): return foo(x)\n" +
                                   "expr = foo(1)");

python/testSrc/com/jetbrains/python/PythonCompletionTest.java

@@ -469,4 +469,9 @@ public class PythonCompletionTest extends PyTestCase {
   public void testCompleteBeforeSyntaxError() { // PY-3792
     doTest();
   }
+
+  // PY-4279
+  public void testFieldReassignment() {
+    doTest();
+  }
 }