Fixed resolve scope for unresolved augmented assignments

python/src/com/jetbrains/python/codeInsight/dataflow/scope/impl/ScopeImpl.java

@@ -14,6 +14,7 @@ import com.jetbrains.python.codeInsight.dataflow.PyReachingDefsSemilattice;
 import com.jetbrains.python.codeInsight.dataflow.scope.Scope;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeVariable;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyAugAssignmentStatementNavigator;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -31,6 +32,7 @@ public class ScopeImpl implements Scope {
   private final ScopeOwner myFlowOwner;
   private volatile Map<String, PsiNamedElement> myNamedElements;
   private volatile List<NameDefiner> myNameDefiners;  // declarations which declare unknown set of names, such as 'from ... import *'
+  private volatile Set<String> myAugAssignments;
 
   public ScopeImpl(final ScopeOwner flowOwner) {
     myFlowOwner = flowOwner;
@@ -87,6 +89,13 @@ public class ScopeImpl implements Scope {
     return myNonlocals.contains(name);
   }
 
+  private boolean isAugAssignment(final String name) {
+    if (myAugAssignments == null || myNestedScopes == null) {
+      collectDeclarations();
+    }
+    return myAugAssignments.contains(name);
+  }
+
   public boolean containsDeclaration(final String name) {
     if (myNamedElements == null || myNameDefiners == null) {
       collectDeclarations();
@@ -97,6 +106,9 @@ public class ScopeImpl implements Scope {
     if (getNamedElement(name) != null) {
       return true;
     }
+    if (isAugAssignment(name)) {
+      return true;
+    }
     for (NameDefiner definer : getNameDefiners()) {
       if (definer.getElementNamed(name) != null) {
         return true;
@@ -150,6 +162,7 @@ public class ScopeImpl implements Scope {
     final List<Scope> nestedScopes = new ArrayList<Scope>();
     final Set<String> globals = new HashSet<String>();
     final Set<String> nonlocals = new HashSet<String>();
+    final Set<String> augAssignments = new HashSet<String>();
     myFlowOwner.acceptChildren(new PyRecursiveElementVisitor() {
       @Override
       public void visitPyTargetExpression(PyTargetExpression node) {
@@ -159,6 +172,14 @@ public class ScopeImpl implements Scope {
         }
       }
 
+      @Override
+      public void visitPyReferenceExpression(PyReferenceExpression node) {
+        if (PyAugAssignmentStatementNavigator.getStatementByTarget(node) != null) {
+          augAssignments.add(node.getName());
+        }
+        super.visitPyReferenceExpression(node);
+      }
+
       @Override
       public void visitPyGlobalStatement(PyGlobalStatement node) {
         for (PyTargetExpression expression : node.getGlobals()) {
@@ -226,5 +247,6 @@ public class ScopeImpl implements Scope {
     myNestedScopes = nestedScopes;
     myGlobals = globals;
     myNonlocals = nonlocals;
+    myAugAssignments = augAssignments;
   }
 }

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

@@ -34,7 +34,10 @@ import com.jetbrains.python.refactoring.PyDefUseUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.*;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.ListIterator;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -261,6 +264,12 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
               uexpr = null;
             }
           }
+          else if (originalOwner != null && ScopeUtil.getResolveScopeOwner(uexpr) != originalOwner && !scope.isGlobal(referencedName)) {
+            final Scope originalScope = ControlFlowCache.getScope(originalOwner);
+            if (originalScope.containsDeclaration(referencedName)) {
+              uexpr = null;
+            }
+          }
         }
       }
       // sort what we got

python/testData/inspections/PyUnresolvedReferencesInspection/augAssignmentDefinedInOuterScope.py

@@ -0,0 +1,3 @@
+x = 1
+def f():
+    <error descr="Unresolved reference 'x'">x</error> += 1

python/testSrc/com/jetbrains/python/inspections/PyUnresolvedReferencesInspectionTest.java

@@ -226,6 +226,11 @@ public class PyUnresolvedReferencesInspectionTest extends PyTestCase {
     doTest();
   }
 
+  // PY-6617
+  public void testAugAssignmentDefinedInOuterScope() {
+    doTest();
+  }
+
   private void doTest() {
     myFixture.configureByFile(TEST_DIRECTORY + getTestName(true) + ".py");
     myFixture.enableInspections(PyUnresolvedReferencesInspection.class);