Test import insertion for type annotations

python/src/com/jetbrains/python/debugger/PySignatureCacheManager.java

@@ -15,22 +15,44 @@
  */
 package com.jetbrains.python.debugger;
 
+import com.google.common.collect.Lists;
 import com.intellij.openapi.components.ServiceManager;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.text.StringUtil;
 import com.jetbrains.python.psi.PyFunction;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.List;
+import java.util.stream.Collectors;
+
 /**
  * @author traff
  */
 public abstract class PySignatureCacheManager {
 
 
+  static final String RETURN_TYPE = "<RETURN_TYPE>";
+
   public static PySignatureCacheManager getInstance(Project project) {
     return ServiceManager.getService(project, PySignatureCacheManager.class);
   }
 
+  public static String signatureToString(PySignature signature) {
+    return signature.getFunctionName() + "\t" + StringUtil.join(arguments(signature), "\t") +
+           (signature.getReturnType() != null
+            ? "\t" + StringUtil.join(
+             signature.getReturnType().getTypesList().stream().map(s -> RETURN_TYPE + ":" + s).collect(Collectors.toList()), "\t") : "");
+  }
+
+  private static List<String> arguments(PySignature signature) {
+    List<String> res = Lists.newArrayList();
+    for (PySignature.NamedParameter param : signature.getArgs()) {
+      res.add(param.getName() + ":" + param.getTypeQualifiedName());
+    }
+    return res;
+  }
+
   public abstract void recordSignature(@NotNull PySignature signature);
 
   @Nullable

python/src/com/jetbrains/python/debugger/PySignatureCacheManagerImpl.java

@@ -18,7 +18,6 @@ package com.jetbrains.python.debugger;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
-import com.google.common.collect.Lists;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.project.Project;
@@ -40,10 +39,8 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.io.IOException;
-import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 
 /**
  * @author traff
@@ -54,7 +51,6 @@ public class PySignatureCacheManagerImpl extends PySignatureCacheManager {
   private final static boolean SHOULD_OVERWRITE_TYPES = false;
 
   public static final FileAttribute CALL_SIGNATURES_ATTRIBUTE = new FileAttribute("call.signatures.attribute", 1, true);
-  private static final String RETURN_TYPE = "<RETURN_TYPE>";
 
   private final Project myProject;
 
@@ -145,21 +141,6 @@ public class PySignatureCacheManagerImpl extends PySignatureCacheManager {
     }
   }
 
-  static String signatureToString(PySignature signature) {
-    return signature.getFunctionName() + "\t" + StringUtil.join(arguments(signature), "\t") +
-           (signature.getReturnType() != null
-            ? "\t" + StringUtil.join(
-             signature.getReturnType().getTypesList().stream().map(s -> RETURN_TYPE + ":" + s).collect(Collectors.toList()), "\t") : "");
-  }
-
-  private static List<String> arguments(PySignature signature) {
-    List<String> res = Lists.newArrayList();
-    for (PySignature.NamedParameter param : signature.getArgs()) {
-      res.add(param.getName() + ":" + param.getTypeQualifiedName());
-    }
-    return res;
-  }
-
   @Nullable
   public String findParameterType(@NotNull PyFunction function, @NotNull String name) {
     final PySignature signature = findSignature(function);

python/testData/intentions/PyAnnotateTypesIntentionTest/importDict.py

@@ -0,0 +1,2 @@
+def get_dict<caret>():
+    return {1: '1'}

python/testData/intentions/PyAnnotateTypesIntentionTest/importDict_after.py

@@ -0,0 +1,4 @@
+from typing import Dict
+
+def get_dict() -> Dict[int, str]:
+    return {1: '1'}

python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java

@@ -15,10 +15,16 @@
  */
 package com.jetbrains.python.intentions;
 
+import com.intellij.codeInsight.intention.IntentionAction;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.debugger.PySignature;
+import com.jetbrains.python.debugger.PySignatureCacheManager;
+import com.jetbrains.python.debugger.PySignatureCacheManagerImpl;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
 
+import java.io.IOException;
+
 /**
  * @author traff
  */
@@ -45,6 +51,29 @@ public class PyAnnotateTypesIntentionTest extends PyIntentionTestCase {
   public void testTypeComment() {
     doTest(PyBundle.message("INTN.annotate.types"), LanguageLevel.PYTHON27);
   }
+
+  public void testImportDict() throws IOException {
+    PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), LanguageLevel.PYTHON30);
+    try {
+      final String testFileName = getTestName(true);
+      myFixture.configureByFile(testFileName + ".py");
+
+      String signature = PySignatureCacheManager.signatureToString(
+        new PySignature(myFixture.getFile().getVirtualFile().getCanonicalPath(), "get_dict").addReturnType("Dict[int, str]"));
+      PySignatureCacheManagerImpl.CALL_SIGNATURES_ATTRIBUTE.writeAttributeBytes(myFixture.getFile().getVirtualFile(),
+                                                                                signature.getBytes());
+
+      final IntentionAction intentionAction = myFixture.findSingleIntention(PyBundle.message("INTN.annotate.types"));
+      assertNotNull(intentionAction);
+      myFixture.launchAction(intentionAction);
+      myFixture.checkResultByFile(testFileName + "_after.py", true);
+    }
+    finally {
+      PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null);
+    }
+
+    doTest();
+  }
   
   
   private void doTest() {