Merge remote-tracking branch 'origin/master'

python/src/com/jetbrains/python/codeInsight/dataflow/scope/ScopeUtil.java

@@ -5,7 +5,6 @@ import com.intellij.codeInsight.controlflow.Instruction;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.StubBasedPsiElement;
 import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 import com.jetbrains.python.codeInsight.controlflow.ReadWriteInstruction;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
@@ -19,6 +18,9 @@ import org.jetbrains.annotations.Nullable;
 import java.util.ArrayList;
 import java.util.Collection;
 
+import static com.intellij.psi.util.PsiTreeUtil.getParentOfType;
+import static com.intellij.psi.util.PsiTreeUtil.isAncestor;
+
 /**
  * @author oleg
  */
@@ -29,7 +31,7 @@ public class ScopeUtil {
   @Nullable
   public static PsiElement getParameterScope(final PsiElement element){
     if (element instanceof PyNamedParameter){
-      final PyFunction function = PsiTreeUtil.getParentOfType(element, PyFunction.class, false);
+      final PyFunction function = getParentOfType(element, PyFunction.class, false);
       if (function != null){
         return function;
       }
@@ -52,8 +54,13 @@ public class ScopeUtil {
     return null;
   }
 
+  /**
+   * Return the scope owner for the element.
+   *
+   * Scope owner is not always the first ScopeOwner parent of the element. Some elements are resolved in outer scopes.
+   */
   @Nullable
-  public static ScopeOwner getScopeOwner(PsiElement element) {
+  public static ScopeOwner getScopeOwner(@Nullable PsiElement element) {
     if (element instanceof StubBasedPsiElement) {
       final StubElement stub = ((StubBasedPsiElement)element).getStub();
       if (stub != null) {
@@ -65,24 +72,31 @@ public class ScopeUtil {
           }
           parentStub = parentStub.getParentStub();
         }
+        return null;
       }
     }
-    return PsiTreeUtil.getParentOfType(element, ScopeOwner.class);
-  }
-
-  @Nullable
-  public static ScopeOwner getResolveScopeOwner(@NotNull PsiElement element) {
-    final ScopeOwner firstOwner = getScopeOwner(element);
+    final ScopeOwner firstOwner = getParentOfType(element, ScopeOwner.class);
     if (firstOwner == null) {
       return null;
     }
-    final ScopeOwner nextOwner = getScopeOwner(firstOwner);
-    final PyElement decoratorOrParameterAncestor = PsiTreeUtil.getParentOfType(element, PyDecorator.class, PyParameter.class);
-    if (decoratorOrParameterAncestor != null && !PsiTreeUtil.isAncestor(decoratorOrParameterAncestor, firstOwner, true)) {
+    final ScopeOwner nextOwner = getParentOfType(firstOwner, ScopeOwner.class);
+    // References in decorator expressions are resolved outside of the function (if the lambda is not inside the decorator)
+    final PyElement decoratorAncestor = getParentOfType(element, PyDecorator.class);
+    if (decoratorAncestor != null && !isAncestor(decoratorAncestor, firstOwner, true)) {
       return nextOwner;
     }
-    final PyClass containingClass = PsiTreeUtil.getParentOfType(element, PyClass.class);
-    if (containingClass != null && PsiTreeUtil.isAncestor(containingClass.getSuperClassExpressionList(), element, false)) {
+    // References in default values of parameters are resolved outside of the function (if the lambda is not inside the default value)
+    final PyParameter parameterAncestor = getParentOfType(element, PyParameter.class);
+    if (parameterAncestor != null && !isAncestor(parameterAncestor, firstOwner, true)) {
+      final PyExpression defaultValue = parameterAncestor.getDefaultValue();
+      if (element != null && isAncestor(defaultValue, element, false)) {
+        return nextOwner;
+      }
+    }
+    // Superclasses are resolved outside of the class
+    final PyClass containingClass = getParentOfType(element, PyClass.class);
+    if (containingClass != null && element != null &&
+        isAncestor(containingClass.getSuperClassExpressionList(), element, false)) {
       return nextOwner;
     }
     return firstOwner;
@@ -92,10 +106,6 @@ public class ScopeUtil {
   public static ScopeOwner getDeclarationScopeOwner(PsiElement anchor, String name) {
     PsiElement element = anchor;
     if (name != null) {
-      // References in default values of parameters are defined somewhere in outer scopes, as well as references in decorators
-      if (PsiTreeUtil.getParentOfType(anchor, PyParameter.class, PyDecorator.class) != null) {
-        element = getScopeOwner(anchor);
-      }
       final ScopeOwner originalScopeOwner = getScopeOwner(element);
       ScopeOwner scopeOwner = originalScopeOwner;
       while (scopeOwner != null) {

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;
 
 /**
@@ -181,7 +184,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;
   }
 
   /**
@@ -223,7 +240,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
     PsiElement uexpr = processor.getResult();
     if (uexpr != null) {
       if (processor.getDefiners().isEmpty()) {
-        final ScopeOwner originalOwner = ScopeUtil.getResolveScopeOwner(realContext);
+        final ScopeOwner originalOwner = ScopeUtil.getScopeOwner(realContext);
         final ScopeOwner owner = ScopeUtil.getScopeOwner(uexpr);
         if (owner != null) {
           final Scope scope = ControlFlowCache.getScope(owner);
@@ -233,19 +250,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) {
@@ -260,6 +264,12 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
               uexpr = null;
             }
           }
+          else if (owner != originalOwner && originalOwner != null && !scope.isGlobal(referencedName)) {
+            final Scope originalScope = ControlFlowCache.getScope(originalOwner);
+            if (originalScope.containsDeclaration(referencedName)) {
+              uexpr = null;
+            }
+          }
         }
       }
       // sort what we got
@@ -360,7 +370,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/psi/resolve/PyResolveUtil.java

@@ -85,7 +85,7 @@ public class PyResolveUtil {
                                   @Nullable PsiElement roof) {
     // Use real context here to enable correct completion and resolve in case of PyExpressionCodeFragment!!!
     final PsiElement realContext = PyPsiUtils.getRealContext(element);
-    final ScopeOwner originalOwner = ScopeUtil.getResolveScopeOwner(realContext);
+    final ScopeOwner originalOwner = ScopeUtil.getScopeOwner(realContext);
     final PsiElement parent = element.getParent();
     final boolean isGlobalOrNonlocal = parent instanceof PyGlobalStatement || parent instanceof PyNonlocalStatement;
     ScopeOwner owner = originalOwner;

python/src/com/jetbrains/python/refactoring/move/PyMoveClassOrFunctionProcessor.java

@@ -135,7 +135,7 @@ public class PyMoveClassOrFunctionProcessor extends BaseRefactoringProcessor {
     List<PsiElement> topLevelAtDestination = new ArrayList<PsiElement>();
     for (UsageInfo usage : usages) {
       final PsiElement e = usage.getElement();
-      if (e != null && ScopeUtil.getResolveScopeOwner(e) == destination && getImportStatementByElement(e) == null) {
+      if (e != null && ScopeUtil.getScopeOwner(e) == destination && getImportStatementByElement(e) == null) {
         PsiElement topLevel = PsiTreeUtil.findFirstParent(e, new Condition<PsiElement>() {
           @Override
           public boolean value(PsiElement element) {

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/inspections/PyUnresolvedReferencesInspection/augAssignmentDefinedInOuterScope.py

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

python/testData/override/kwargs.py

@@ -1,6 +1,6 @@
-class A():
-    def m(self, *args, **kwargs):
-        pass
-
-class B(A):
-    pass
+class A():
+    def m(self, *args, **kwargs):
+        pass
+
+class B(A):
+    pass

python/testData/override/kwargs_after.py

@@ -1,7 +1,7 @@
-class A():
-    def m(self, *args, **kwargs):
-        pass
-
-class B(A):
-    def m(self, *args, **kwargs):
-        <selection>super().m(*args, **kwargs)</selection>
+class A():
+    def m(self, *args, **kwargs):
+        pass
+
+class B(A):
+    def m(self, *args, **kwargs):
+        <selection>super().m(*args, **kwargs)</selection>

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);
+  }
 }

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);