convert RatedResolveResult to class from interface; initial implementation of unused import highlighting

python/src/com/jetbrains/python/inspections/PyUnresolvedReferencesInspection.java

@@ -23,6 +23,7 @@ import com.jetbrains.python.psi.patterns.Matcher;
 import com.jetbrains.python.psi.patterns.ParentMatcher;
 import com.jetbrains.python.psi.patterns.SyntaxMatchers;
 import com.jetbrains.python.psi.resolve.CollectProcessor;
+import com.jetbrains.python.psi.resolve.ImportedResolveResult;
 import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import com.jetbrains.python.psi.resolve.ResolveImportUtil;
 import com.jetbrains.python.psi.stubs.PyClassNameIndex;
@@ -39,6 +40,8 @@ import java.util.*;
  * Date: Nov 15, 2008
  */
 public class PyUnresolvedReferencesInspection extends LocalInspectionTool {
+  private final ThreadLocal<Visitor> myLastVisitor = new ThreadLocal<Visitor>();
+
   @Nls
   @NotNull
   public String getGroupDisplayName() {
@@ -81,10 +84,22 @@ public class PyUnresolvedReferencesInspection extends LocalInspectionTool {
   @NotNull
   @Override
   public PsiElementVisitor buildVisitor(@NotNull final ProblemsHolder holder, final boolean isOnTheFly) {
-    return new Visitor(holder);
+    final Visitor visitor = new Visitor(holder);
+    myLastVisitor.set(visitor);
+    return visitor;
+  }
+
+  @Override
+  public void inspectionFinished(LocalInspectionToolSession session) {
+    final Visitor visitor = myLastVisitor.get();
+    assert visitor != null;
+    visitor.highlightUnusedImports();
+    myLastVisitor.remove();
   }
 
   public static class Visitor extends PyInspectionVisitor {
+    private Set<NameDefiner> myUsedImports = Collections.synchronizedSet(new HashSet<NameDefiner>());
+    private Set<PyImportElement> myAllImports = Collections.synchronizedSet(new HashSet<PyImportElement>());
 
     public Visitor(final ProblemsHolder holder) {
       super(holder);
@@ -256,6 +271,12 @@ public class PyUnresolvedReferencesInspection extends LocalInspectionTool {
       }
     }
 
+    @Override
+    public void visitPyImportElement(PyImportElement node) {
+      super.visitPyImportElement(node);
+      myAllImports.add(node);
+    }
+
     @Override
     public void visitPyElement(final PyElement node) {
       super.visitPyElement(node);
@@ -269,13 +290,23 @@ public class PyUnresolvedReferencesInspection extends LocalInspectionTool {
         boolean unresolved;
         if (reference instanceof PsiPolyVariantReference) {
           final PsiPolyVariantReference poly = (PsiPolyVariantReference)reference;
-          unresolved = (poly.multiResolve(false).length == 0);
+          final ResolveResult[] resolveResults = poly.multiResolve(false);
+          unresolved = (resolveResults.length == 0);
+          for (ResolveResult resolveResult : resolveResults) {
+            if (resolveResult instanceof ImportedResolveResult) {
+              myUsedImports.addAll(((ImportedResolveResult)resolveResult).getNameDefiners());
+            }
+          }
         }
         else {
           unresolved = (reference.resolve() == null);
         }
         if (unresolved) {
           registerUnresolvedReferenceProblem(node, reference, severity);
+          // don't highlight unresolved imports as unused
+          if (node.getParent() instanceof PyImportElement) {
+            myAllImports.remove(node.getParent());
+          }
         }
       }
     }
@@ -409,5 +440,34 @@ public class PyUnresolvedReferencesInspection extends LocalInspectionTool {
         });
       }
     }
+
+    public void highlightUnusedImports() {
+      myAllImports.removeAll(myUsedImports);
+      final NameDefiner[] unusedImports = myAllImports.toArray(new NameDefiner[myAllImports.size()]);
+
+      Set<PyImportStatementBase> unusedStatements = new HashSet<PyImportStatementBase>();
+      for (NameDefiner unusedImport : unusedImports) {
+        PyImportStatementBase importStatement = PsiTreeUtil.getParentOfType(unusedImport, PyImportStatementBase.class);
+        if (importStatement != null && !unusedStatements.contains(importStatement)) {
+          if (areAllImportsUnused(importStatement)) {
+            unusedStatements.add(importStatement);
+            registerProblem(importStatement, "Unused import statement", ProblemHighlightType.LIKE_UNUSED_SYMBOL, null);
+          }
+          else {
+            registerProblem(unusedImport, "Unused import statement", ProblemHighlightType.LIKE_UNUSED_SYMBOL, null);
+          }
+        }
+      }
+    }
+
+    private boolean areAllImportsUnused(PyImportStatementBase importStatement) {
+      final PyImportElement[] elements = importStatement.getImportElements();
+      for (PyImportElement element : elements) {
+        if (!myAllImports.contains(element)) {
+          return false;
+        }
+      }
+      return true;
+    }
   }
 }

python/src/com/jetbrains/python/psi/PyElementVisitor.java

@@ -1,7 +1,6 @@
 package com.jetbrains.python.psi;
 
 import com.intellij.psi.PsiElementVisitor;
-import com.jetbrains.python.psi.impl.PySubscriptionExpressionImpl;
 
 /**
  * Visitor for python-specific nodes.
@@ -214,4 +213,8 @@ public class PyElementVisitor extends PsiElementVisitor {
   public void visitPySubscriptionExpression(PySubscriptionExpression node) {
     visitPyExpression(node);
   }
+
+  public void visitPyImportElement(PyImportElement node) {
+    visitPyElement(node);
+  }
 }

python/src/com/jetbrains/python/psi/PyFromImportStatement.java

@@ -8,7 +8,7 @@ import org.jetbrains.annotations.Nullable;
 /**
  * Describes "from ... import" statements.
  */
-public interface PyFromImportStatement extends PyStatement, StubBasedPsiElement<PyFromImportStatementStub> {
+public interface PyFromImportStatement extends PyImportStatementBase, StubBasedPsiElement<PyFromImportStatementStub> {
   boolean isStarImport();
 
   /**
@@ -20,11 +20,6 @@ public interface PyFromImportStatement extends PyStatement, StubBasedPsiElement<
   @Nullable
   PyQualifiedName getImportSourceQName();
 
-  /**
-   * @return elements that constitute the "import" clause
-   */
-  PyImportElement[] getImportElements();
-
   /**
    * @return the star in "from ... import *"
    */

python/src/com/jetbrains/python/psi/PyImportStatement.java

@@ -6,6 +6,5 @@ import com.jetbrains.python.psi.stubs.PyImportStatementStub;
 /**
  * @author yole
  */
-public interface PyImportStatement extends PyStatement, StubBasedPsiElement<PyImportStatementStub> {
-  PyImportElement[] getImportElements();
+public interface PyImportStatement extends PyImportStatementBase, StubBasedPsiElement<PyImportStatementStub> {
 }

python/src/com/jetbrains/python/psi/PyImportStatementBase.java

@@ -0,0 +1,11 @@
+package com.jetbrains.python.psi;
+
+/**
+ * @author yole
+ */
+public interface PyImportStatementBase extends PyStatement {
+  /**
+   * @return elements that constitute the "import" clause
+   */
+  PyImportElement[] getImportElements();
+}

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

@@ -226,4 +226,9 @@ public class PyImportElementImpl extends PyBaseElementImpl<PyImportElementStub>
   public boolean mustResolveOutside() {
     return true; // formally
   }
+
+  @Override
+  protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
+    pyVisitor.visitPyImportElement(this);
+  }
 }

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

@@ -103,28 +103,14 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
           PyFunction init = cls.findMethodByName(PyNames.INIT, false);
           if (init != null) {
             // replace
-            final PyFunction the_init = init;
-            it.set(new RatedResolveResult(){
-              public int getRate() { return rrr.getRate(); }
-
-              public PsiElement getElement() { return the_init; }
-
-              public boolean isValidResult() { return true; }
-            });
+            it.set(rrr.replace(init));
           }
           else { // init not found; maybe it's ancestor's
             for (PyClass ancestor : cls.iterateAncestors()) {
               init = ancestor.findMethodByName(PyNames.INIT, false);
               if (init != null) {
-                final PyFunction the_init = init;
                 // add to resuls as low priority
-                it.add(new RatedResolveResult(){
-                  public int getRate() { return RATE_LOW; }
-
-                  public PsiElement getElement() { return the_init; }
-
-                  public boolean isValidResult() { return true; }
-                });
+                it.add(new RatedResolveResult(RatedResolveResult.RATE_LOW, init));
                 break;
               }
             }
@@ -148,19 +134,9 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
     // Allows to add non-null elements and discard nulls in a hassle-free way.
     public boolean poke(final PsiElement what, final int rate) {
       if (what == null) return false;
-      super.add(new RatedResolveResult() {
-        public int getRate() { return rate; }
-
-        public PsiElement getElement() { return what; }
-
-        public boolean isValidResult() { return true; }
-      });
+      super.add(new RatedResolveResult(rate, what));
       return true;
     }
-
-    public void pokeAll(Collection<PsiElement> elts, int rate) {
-      for (PsiElement elt : elts) poke(elt, rate);
-    }
   }
 
   /**
@@ -213,7 +189,9 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
       uexpr = PyResolveUtil.scanOuterContext(new ResolveProcessor(referencedName), realContext);
     }
     uexpr = PyUtil.turnDirIntoInit(uexpr); // treeCrawlUp might have found a dir
-    if (uexpr != null) ret.poke(uexpr, getRate(uexpr));
+    if (uexpr != null) {
+      ret.add(new ImportedResolveResult(uexpr, getRate(uexpr), processor.getDefiners()));
+    }
     return ret;
   }
 

python/src/com/jetbrains/python/psi/resolve/ImplicitResolveResult.java

@@ -6,23 +6,8 @@ import org.jetbrains.annotations.Nullable;
 /**
  * @author yole
  */
-public class ImplicitResolveResult implements RatedResolveResult {
-  @Nullable private final PsiElement myElement;
-
+public class ImplicitResolveResult extends RatedResolveResult {
   public ImplicitResolveResult(@Nullable final PsiElement element) {
-    myElement = element;
-  }
-
-  @Nullable
-  public PsiElement getElement() {
-    return myElement;
-  }
-
-  public boolean isValidResult() {
-    return false;
-  }
-
-  public int getRate() {
-    return RATE_LOW;
+    super(RATE_LOW, element);
   }
 }

python/src/com/jetbrains/python/psi/resolve/ImportedResolveResult.java

@@ -0,0 +1,27 @@
+package com.jetbrains.python.psi.resolve;
+
+import com.intellij.psi.PsiElement;
+import com.jetbrains.python.psi.NameDefiner;
+
+import java.util.List;
+
+/**
+ * @author yole
+ */
+public class ImportedResolveResult extends RatedResolveResult {
+  private final List<NameDefiner> myNameDefiners;
+
+  public ImportedResolveResult(PsiElement element, int rate, List<NameDefiner> nameDefiners) {
+    super(rate, element);
+    myNameDefiners = nameDefiners;
+  }
+
+  public List<NameDefiner> getNameDefiners() {
+    return myNameDefiners;
+  }
+
+  @Override
+  public RatedResolveResult replace(PsiElement what) {
+    return new ImportedResolveResult(what, getRate(), myNameDefiners);
+  }
+}

python/src/com/jetbrains/python/psi/resolve/RatedResolveResult.java

@@ -1,5 +1,6 @@
 package com.jetbrains.python.psi.resolve;
 
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.ResolveResult;
 
 /**
@@ -7,29 +8,50 @@ import com.intellij.psi.ResolveResult;
  * User: dcheryasov
  * Date: Dec 5, 2008 11:06:30 AM
  */
-public interface RatedResolveResult extends ResolveResult {
+public class RatedResolveResult implements ResolveResult {
+  private final int myRate;
+  private final PsiElement myWhat;
 
-  RatedResolveResult[] EMPTY_ARRAY = new RatedResolveResult[0];
+  public RatedResolveResult(int rate, PsiElement what) {
+    myRate = rate;
+    myWhat = what;
+  }
+
+  public static final RatedResolveResult[] EMPTY_ARRAY = new RatedResolveResult[0];
+
+  public PsiElement getElement() {
+    return myWhat;
+  }
+
+  public boolean isValidResult() {
+    return true;
+  }
 
   /**
    * Results with higher rate are shown higher in the list of multiResolve().
    * @see com.intellij.psi.PsiPolyVariantReference#multiResolve(boolean)
    * @return desired rate. If in doubt, use 0.
    */
-  int getRate();
+  public int getRate() {
+    return myRate;
+  }
+
+  public RatedResolveResult replace(PsiElement what) {
+    return new RatedResolveResult(myRate, what);
+  }
 
   /**
    * For unusual items that need to be on top.
    */
-  int RATE_HIGH = 1000;
+  public static final int RATE_HIGH = 1000;
 
   /**
    * For regular references.
    */
-  int RATE_NORMAL = 0;
+  public static final int RATE_NORMAL = 0;
 
   /**
    * For additional, less important results.
    */
-  int RATE_LOW = -1000;
+  public static final int RATE_LOW = -1000;
 }