PY-18521 Consider references inside function type comments to mark corresponding imports as used

Additionally PyImportOptimizer visits PSI comments collecting
information about unused imports, otherwise it doesn't take into
account references inside any kind of type comments.

python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java

@@ -19,6 +19,7 @@ import com.google.common.collect.Ordering;
 import com.intellij.codeInspection.LocalInspectionToolSession;
 import com.intellij.lang.ImportOptimizer;
 import com.intellij.openapi.util.Condition;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.codeInsight.imports.AddImportHelper.ImportPriority;
@@ -49,8 +50,8 @@ public class PyImportOptimizer implements ImportOptimizer {
                                                                                                           Collections.<String>emptyList());
     file.accept(new PyRecursiveElementVisitor() {
       @Override
-      public void visitPyElement(PyElement node) {
-        super.visitPyElement(node);
+      public void visitElement(PsiElement node) {
+        super.visitElement(node);
         node.accept(visitor);
       }
     });

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

@@ -42,6 +42,7 @@ import com.jetbrains.python.PyCustomType;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
 import com.jetbrains.python.codeInsight.PyCustomMember;
+import com.jetbrains.python.codeInsight.PyFunctionTypeCommentReferenceContributor;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.codeInsight.imports.AutoImportHintAction;
@@ -286,6 +287,13 @@ public class PyUnresolvedReferencesInspection extends PyInspection {
       if (comment instanceof PsiLanguageInjectionHost) {
         processInjection((PsiLanguageInjectionHost)comment);
       }
+      if (PyFunctionTypeCommentReferenceContributor.TYPE_COMMENT_PATTERN.accepts(comment)) {
+        for (PsiReference reference : comment.getReferences()) {
+          if (reference instanceof PsiPolyVariantReference) {
+            markTargetImportsAsUsed((PsiPolyVariantReference)reference);
+          }
+        }
+      }
     }
 
     @Override
@@ -323,15 +331,7 @@ public class PyUnresolvedReferencesInspection extends PyInspection {
               if (element instanceof PyReferenceOwner) {
                 final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(myTypeEvalContext);
                 final PsiPolyVariantReference reference = ((PyReferenceOwner)element).getReference(resolveContext);
-                final ResolveResult[] resolveResults = reference.multiResolve(false);
-                for (ResolveResult resolveResult : resolveResults) {
-                  if (resolveResult instanceof ImportedResolveResult) {
-                    final PyImportedNameDefiner definer = ((ImportedResolveResult)resolveResult).getDefiner();
-                    if (definer != null) {
-                      myUsedImports.add(definer);
-                    }
-                  }
-                }
+                markTargetImportsAsUsed(reference);
               }
             }
           }.visitElement(pair.getFirst());
@@ -339,6 +339,18 @@ public class PyUnresolvedReferencesInspection extends PyInspection {
       }
     }
 
+    private void markTargetImportsAsUsed(@NotNull PsiPolyVariantReference reference) {
+      final ResolveResult[] resolveResults = reference.multiResolve(false);
+      for (ResolveResult resolveResult : resolveResults) {
+        if (resolveResult instanceof ImportedResolveResult) {
+          final PyImportedNameDefiner definer = ((ImportedResolveResult)resolveResult).getDefiner();
+          if (definer != null) {
+            myUsedImports.add(definer);
+          }
+        }
+      }
+    }
+
     private void processReference(PyElement node, @Nullable PsiReference reference) {
       if (!isEnabled(node) || reference == null || reference.isSoft()) {
         return;

python/testData/inspections/PyUnresolvedReferencesInspection/functionTypeCommentUsesImportsFromTyping.py

@@ -0,0 +1,7 @@
+from typing import List, <warning descr="Unused import statement">Optional</warning>
+
+
+def f(x, y):
+    # type: (int, List[int]) -> str
+    y.append(x)
+    return 'foo'

python/testData/optimizeImports/importsFromTypingUnusedInTypeComments.after.py

@@ -0,0 +1,10 @@
+from typing import List, Set
+
+
+def f(x, y):
+    # type: (int, List[int]) -> str
+    y.append(x)
+    return 'foo'
+
+
+xs = {1, 2, 3} # type: Set[int]

python/testData/optimizeImports/importsFromTypingUnusedInTypeComments.py

@@ -0,0 +1,10 @@
+from typing import List, Optional, Set
+
+
+def f(x, y):
+    # type: (int, List[int]) -> str
+    y.append(x)
+    return 'foo'
+
+
+xs = {1, 2, 3} # type: Set[int]

python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java

@@ -81,6 +81,12 @@ public class PyOptimizeImportsTest extends PyTestCase {
     myFixture.checkResultByFile(testName + "/main.after.py");
   }
 
+  // PY-18521
+  public void testImportsFromTypingUnusedInTypeComments() {
+    myFixture.copyDirectoryToProject("../typing", "");
+    doTest();
+  }
+
   private void doTest() {
     myFixture.configureByFile(getTestName(true) + ".py");
     OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));

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

@@ -542,6 +542,12 @@ public class PyUnresolvedReferencesInspectionTest extends PyInspectionTestCase {
     doTest();
   }
 
+  // PY-18521
+  public void testFunctionTypeCommentUsesImportsFromTyping() {
+    myFixture.copyDirectoryToProject("typing", "");
+    doTest();
+  }
+
   @NotNull
   @Override
   protected Class<? extends PyInspection> getInspectionClass() {