From 26db370e5ee9a5323e9fb06606f87baf8ef100ca Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Mon, 24 Oct 2016 15:41:22 +0300
Subject: [PATCH] PY-20043 Unused import in notebook for usages in next cell

---
 .../resources/META-INF/ipython-notebook.xml   |  1 +
 .../psi/IpnbUnresolvedReferenceSkipper.java   | 70 +++++++++++++++++++
 2 files changed, 71 insertions(+)
 create mode 100644 python/ipnb/src/org/jetbrains/plugins/ipnb/psi/IpnbUnresolvedReferenceSkipper.java

diff --git a/python/ipnb/resources/META-INF/ipython-notebook.xml b/python/ipnb/resources/META-INF/ipython-notebook.xml
index ea2ac8ca0acf..1dbc44eb4753 100644
--- a/python/ipnb/resources/META-INF/ipython-notebook.xml
+++ b/python/ipnb/resources/META-INF/ipython-notebook.xml
@@ -19,6 +19,7 @@
     <dialectsTokenSetContributor implementation="org.jetbrains.plugins.ipnb.psi.IpnbPyTokenSetContributor"/>
     <visitorFilter language="IpnbPython" implementationClass="org.jetbrains.plugins.ipnb.IpnbVisitorFilter"/>
     <pep8ProblemSuppressor implementation="org.jetbrains.plugins.ipnb.IpnbPep8ProblemSuppressor"/>
+    <unresolvedReferenceSkipper implementation="org.jetbrains.plugins.ipnb.psi.IpnbUnresolvedReferenceSkipper"/>
   </extensions>
   <actions>
     <action class="org.jetbrains.plugins.ipnb.editor.actions.IpnbRunCellInplaceAction" id="IpnbRunCellInplaceAction"
diff --git a/python/ipnb/src/org/jetbrains/plugins/ipnb/psi/IpnbUnresolvedReferenceSkipper.java b/python/ipnb/src/org/jetbrains/plugins/ipnb/psi/IpnbUnresolvedReferenceSkipper.java
new file mode 100644
index 000000000000..24cbdfedbfe5
--- /dev/null
+++ b/python/ipnb/src/org/jetbrains/plugins/ipnb/psi/IpnbUnresolvedReferenceSkipper.java
@@ -0,0 +1,70 @@
+package org.jetbrains.plugins.ipnb.psi;
+
+import com.intellij.openapi.editor.Editor;
+import com.intellij.psi.*;
+import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferenceSkipperExtPoint;
+import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyImportedNameDefiner;
+import com.jetbrains.python.psi.PyRecursiveElementVisitor;
+import com.jetbrains.python.psi.PyReferenceOwner;
+import com.jetbrains.python.psi.resolve.ImportedResolveResult;
+import com.jetbrains.python.psi.resolve.PyResolveContext;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.plugins.ipnb.editor.panels.IpnbEditablePanel;
+import org.jetbrains.plugins.ipnb.editor.panels.IpnbFilePanel;
+import org.jetbrains.plugins.ipnb.editor.panels.code.IpnbCodePanel;
+
+import java.util.List;
+
+public class IpnbUnresolvedReferenceSkipper implements PyUnresolvedReferenceSkipperExtPoint {
+  @Override
+  public boolean unusedImportShouldBeSkipped(@NotNull final PyImportedNameDefiner importNameDefiner) {
+    final PsiFile file = importNameDefiner.getContainingFile();
+    if (file instanceof IpnbPyFragment) {
+      final IpnbFilePanel panel = ((IpnbPyFragment)file).getFilePanel();
+      final List<IpnbEditablePanel> panels = panel.getIpnbPanels();
+
+      for (IpnbEditablePanel editablePanel : panels) {
+        if (!(editablePanel instanceof IpnbCodePanel)) continue;
+        final Editor editor = ((IpnbCodePanel)editablePanel).getEditor();
+        final IpnbPyFragment psiFile = (IpnbPyFragment)PsiDocumentManager.getInstance(importNameDefiner.getProject()).getPsiFile(editor.getDocument());
+        if (psiFile == null) continue;
+        final MyVisitor visitor = new MyVisitor(importNameDefiner);
+        psiFile.accept(visitor);
+        if (visitor.used) return true;
+      }
+    }
+    return false;
+  }
+
+  private static class MyVisitor extends PyRecursiveElementVisitor {
+    private final PyImportedNameDefiner myImportNameDefiner;
+    boolean used = false;
+
+    public MyVisitor(PyImportedNameDefiner importNameDefiner) {
+      myImportNameDefiner = importNameDefiner;
+    }
+
+    @Override
+    public void visitPyElement(PyElement node) {
+      super.visitPyElement(node);
+      if (node instanceof PyReferenceOwner) {
+        final PyResolveContext resolveContext = PyResolveContext.noImplicits();
+        final PsiPolyVariantReference reference = ((PyReferenceOwner)node).getReference(resolveContext);
+        PsiElement target = null;
+        final ResolveResult[] resolveResults = reference.multiResolve(false);
+        for (ResolveResult resolveResult : resolveResults) {
+          if (target == null && resolveResult.isValidResult()) {
+            target = resolveResult.getElement();
+          }
+          if (resolveResult instanceof ImportedResolveResult) {
+            final PyImportedNameDefiner definer = ((ImportedResolveResult)resolveResult).getDefiner();
+            if (myImportNameDefiner.equals(definer)) {
+              used = true;
+            }
+          }
+        }
+      }
+    }
+  }
+}