variations on the theme of Find Usages for local variables (PY-527)

python/src/com/jetbrains/python/findUsages/PyFindUsagesHandlerFactory.java

@@ -4,6 +4,8 @@ import com.intellij.find.findUsages.FindUsagesHandler;
 import com.intellij.find.findUsages.FindUsagesHandlerFactory;
 import com.intellij.psi.PsiElement;
 import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.PyTargetExpression;
+import com.jetbrains.python.psi.impl.PyTargetExpressionImpl;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -12,11 +14,19 @@ import org.jetbrains.annotations.NotNull;
 public class PyFindUsagesHandlerFactory extends FindUsagesHandlerFactory {
   @Override
   public boolean canFindUsages(@NotNull PsiElement element) {
-    return element instanceof PyClass;
+    return element instanceof PyClass || isLocal(element);
+  }
+
+  private static boolean isLocal(PsiElement element) {
+    if (!(element instanceof PyTargetExpression)) return false;
+    return ((PyTargetExpressionImpl) element).getRedefinitionScope() != null;
   }
 
   @Override
   public FindUsagesHandler createFindUsagesHandler(@NotNull PsiElement element, boolean forHighlightUsages) {
-    return new PyClassFindUsagesHandler((PyClass) element);
+    if (element instanceof PyClass) {
+      return new PyClassFindUsagesHandler((PyClass)element);
+    }
+    return new PyLocalFindUsagesHandler(element);
   }
 }

python/src/com/jetbrains/python/findUsages/PyLocalFindUsagesHandler.java

@@ -0,0 +1,49 @@
+package com.jetbrains.python.findUsages;
+
+import com.intellij.find.findUsages.FindUsagesHandler;
+import com.intellij.find.findUsages.FindUsagesOptions;
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.util.Comparing;
+import com.intellij.psi.PsiElement;
+import com.intellij.usageView.UsageInfo;
+import com.intellij.util.Processor;
+import com.jetbrains.python.psi.PyRecursiveElementVisitor;
+import com.jetbrains.python.psi.PyTargetExpression;
+import com.jetbrains.python.psi.impl.PyTargetExpressionImpl;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * @author yole
+ */
+public class PyLocalFindUsagesHandler extends FindUsagesHandler {
+  public PyLocalFindUsagesHandler(@NotNull PsiElement psiElement) {
+    super(psiElement);
+  }
+
+  @Override
+  public void processElementUsages(@NotNull final PsiElement element,
+                                   @NotNull final Processor<UsageInfo> processor,
+                                   @NotNull FindUsagesOptions options) {
+    super.processElementUsages(element, processor, options);
+    if (element instanceof PyTargetExpressionImpl) {
+      final PsiElement redefinitionScope = ((PyTargetExpressionImpl)element).getRedefinitionScope();
+      if (redefinitionScope != null) {
+        ApplicationManager.getApplication().runReadAction(new Runnable() {
+          public void run() {
+            redefinitionScope.accept(new PyRecursiveElementVisitor() {
+              private boolean found = false;
+
+              @Override
+              public void visitPyTargetExpression(PyTargetExpression node) {
+                if (Comparing.equal(node.getName(), ((PyTargetExpression)element).getName()) && !found) {
+                  found = true;
+                  processor.process(new UsageInfo(node));
+                }
+              }
+            });
+          }
+        });
+      }
+    }
+  }
+}

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

@@ -364,7 +364,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
   }
 
   public boolean isSoft() {
-    return false;
+    return myElement instanceof PyTargetExpression;
   }
 
   public HighlightSeverity getUnresolvedHighlightSeverity() {

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

@@ -176,12 +176,25 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
   }
 
   public boolean isRedefiningAssignment() {
-    PyFunction containingFunction = PsiTreeUtil.getParentOfType(this, PyFunction.class);
-    final PyRedefinitionVisitor visitor = new PyRedefinitionVisitor();
-    if (containingFunction != null) {
-      containingFunction.acceptChildren(visitor);
+    PsiElement redefinitionScope = getRedefinitionScope();
+    if (redefinitionScope != null) {
+      final PyRedefinitionVisitor visitor = new PyRedefinitionVisitor();
+      redefinitionScope.acceptChildren(visitor);
+      return visitor.myResult != null;
     }
-    return visitor.myResult != null;
+    return false;
+  }
+
+  @Nullable
+  public PsiElement getRedefinitionScope() {
+    PyFunction containingFunction = PsiTreeUtil.getParentOfType(this, PyFunction.class);
+    if (containingFunction != null) {
+      return containingFunction;
+    }
+    if (PsiTreeUtil.getParentOfType(this, PyClass.class) == null) {
+      return getContainingFile();
+    }
+    return null;
   }
 
   private class PyRedefinitionVisitor extends PyRecursiveElementVisitor {
@@ -204,5 +217,15 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
         myResult = node;
       }
     }
+
+    @Override
+    public void visitPyClass(PyClass node) {
+      // don't go into classes
+    }
+
+    @Override
+    public void visitPyFunction(PyFunction node) {
+      // don't go into functions
+    }
   }
 }

python/testData/findUsages/ReassignedLocalUsages.py

@@ -0,0 +1,3 @@
+a<caret>aa = 1234
+aaa = 123
+aaa

python/testSrc/com/jetbrains/python/PyFindUsagesTest.java

@@ -18,4 +18,9 @@ public class PyFindUsagesTest extends PyLightFixtureTestCase {
     final Collection<UsageInfo> usages = myFixture.testFindUsages("findUsages/ClassUsages.py");
     assertEquals(1, usages.size());
   }
+
+  public void testReassignedLocalUsages() { // PY-527
+    final Collection<UsageInfo> usages = myFixture.testFindUsages("findUsages/ReassignedLocalUsages.py");
+    assertEquals(3, usages.size());
+  }
 }