From 30738e534bc0aa9d9349fbfe4caf88458191a698 Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Fri, 7 Dec 2012 14:27:33 +0400
Subject: [PATCH] Fixed resolve for augmented assignments in cycles and for
 undefined variables (PY-7970)

---
 .../psi/impl/references/PyReferenceImpl.java  | 31 ++++++++++---------
 .../rename/RenamePyVariableProcessor.java     |  3 +-
 .../resolve/AugmentedAfterAugmented.py        |  6 ++++
 .../testData/resolve/AugmentedAssignment.py   |  7 +++++
 .../com/jetbrains/python/PyResolveTest.java   | 18 +++++++++++
 5 files changed, 49 insertions(+), 16 deletions(-)
 create mode 100644 python/testData/resolve/AugmentedAfterAugmented.py
 create mode 100644 python/testData/resolve/AugmentedAssignment.py

diff --git a/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java b/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
index 912bd282e558..cd2b8af5e1de 100644
--- a/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
+++ b/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) {
diff --git a/python/src/com/jetbrains/python/refactoring/rename/RenamePyVariableProcessor.java b/python/src/com/jetbrains/python/refactoring/rename/RenamePyVariableProcessor.java
index 5850ae45cdc3..fb377fb8ca19 100644
--- a/python/src/com/jetbrains/python/refactoring/rename/RenamePyVariableProcessor.java
+++ b/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
diff --git a/python/testData/resolve/AugmentedAfterAugmented.py b/python/testData/resolve/AugmentedAfterAugmented.py
new file mode 100644
index 000000000000..af870d6359d0
--- /dev/null
+++ b/python/testData/resolve/AugmentedAfterAugmented.py
@@ -0,0 +1,6 @@
+foo = 1
+foo += 1
+
+while True:
+    foo += 2
+#   <ref>
diff --git a/python/testData/resolve/AugmentedAssignment.py b/python/testData/resolve/AugmentedAssignment.py
new file mode 100644
index 000000000000..5beb7e1ff17d
--- /dev/null
+++ b/python/testData/resolve/AugmentedAssignment.py
@@ -0,0 +1,7 @@
+foo = 1
+
+while True:
+    foo += 2
+#   <ref>
+
+print(foo)
diff --git a/python/testSrc/com/jetbrains/python/PyResolveTest.java b/python/testSrc/com/jetbrains/python/PyResolveTest.java
index c973512352c3..b79ca31463df 100644
--- a/python/testSrc/com/jetbrains/python/PyResolveTest.java
+++ b/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);
+  }
 }