Fixed resolve for augmented assignments in cycles and for undefined variables (PY-7970)

python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java

@@ -181,7 +181,21 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
         ret.poke(definition, getRate(definition));
       }
     }
-    return ret;
+    final ResolveResultList results = new ResolveResultList();
+    for (RatedResolveResult r : ret) {
+      final PsiElement e = r.getElement();
+      if (e == element) {
+        continue;
+      }
+      if (element instanceof PyTargetExpression && PyPsiUtils.isBefore(element, e)) {
+        continue;
+      }
+      else {
+        results.add(r);
+      }
+    }
+
+    return results;
   }
 
   /**
@@ -233,19 +247,6 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
           else if (owner == originalOwner && !scope.isGlobal(referencedName)) {
             final ResolveResultList latest = resolveToLatestDefs(owner, myElement, referencedName);
             if (!latest.isEmpty()) {
-              if (myElement instanceof PyTargetExpression) {
-                final RatedResolveResult result = latest.get(0);
-                final PsiElement element = result.getElement();
-                if (element instanceof PyTargetExpression) {
-                  if (PyPsiUtils.isBefore(element, myElement)) {
-                    return latest;
-                  }
-                  else {
-                    ret.poke(myElement, getRate(myElement));
-                    return ret;
-                  }
-                }
-              }
               return latest;
             }
             if (owner instanceof PyClass) {
@@ -360,7 +361,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
     if (CythonLanguageDialect.isInsideCythonFile(elt) && elt instanceof CythonIncludeStatement) {
       rate = RatedResolveResult.RATE_LOW;
     }
-    else if (elt instanceof PyImportElement || elt instanceof PyStarImportElement) {
+    else if (elt instanceof PyImportElement || elt instanceof PyStarImportElement || elt instanceof PyReferenceExpression) {
       rate = RatedResolveResult.RATE_LOW;
     }
     else if (elt instanceof PyFile) {

python/src/com/jetbrains/python/refactoring/rename/RenamePyVariableProcessor.java

@@ -3,6 +3,7 @@ package com.jetbrains.python.refactoring.rename;
 import com.intellij.psi.PsiElement;
 import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
 import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyReferenceExpression;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -12,7 +13,7 @@ public class RenamePyVariableProcessor extends RenamePyElementProcessor {
   @Override
   public boolean canProcessElement(@NotNull PsiElement element) {
     // extension ordering in python-plugin-common.xml ensures that classes and functions are handled by their own processors
-    return element instanceof PyElement;
+    return element instanceof PyElement && !(element instanceof PyReferenceExpression);
   }
 
   @Override

python/testData/resolve/AugmentedAfterAugmented.py

@@ -0,0 +1,6 @@
+foo = 1
+foo += 1
+
+while True:
+    foo += 2
+#   <ref>

python/testData/resolve/AugmentedAssignment.py

@@ -0,0 +1,7 @@
+foo = 1
+
+while True:
+    foo += 2
+#   <ref>
+
+print(foo)

python/testSrc/com/jetbrains/python/PyResolveTest.java

@@ -503,4 +503,22 @@ public class PyResolveTest extends PyResolveTestCase {
     assertNotNull(target);
     assertTrue(source == target);
   }
+
+  // PY-7970
+  public void testAugmentedAssignment() {
+    assertResolvesTo(PyTargetExpression.class, "foo");
+  }
+
+  // PY-7970
+  public void testAugmentedAfterAugmented() {
+    final PsiReference ref = findReferenceByMarker();
+    final PsiElement source = ref.getElement();
+    final PsiElement resolved = ref.resolve();
+    assertInstanceOf(resolved, PyReferenceExpression.class);
+    assertNotSame(resolved, source);
+    final PyReferenceExpression res = (PyReferenceExpression)resolved;
+    assertNotNull(res);
+    assertEquals("foo", res.getName());
+    assertInstanceOf(res.getParent(), PyAugAssignmentStatement.class);
+  }
 }