From 034f409a6af6d76f87ab1bfb5c71c8f2b82d5c14 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Sat, 27 Jan 2018 21:45:08 +0300
Subject: [PATCH] PY-16412 Import Type for annotations containing types of
 class objects

---
 .../codeInsight/intentions/PyTypeHintGenerationUtil.java  | 3 +++
 .../annotationClassObjectType.py                          | 5 +++++
 .../annotationClassObjectType_after.py                    | 8 ++++++++
 .../intentions/PyAnnotateVariableTypeIntentionTest.java   | 4 ++++
 4 files changed, 20 insertions(+)
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationClassObjectType.py
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationClassObjectType_after.py

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java
index e0bba27a3522..992219572ede 100644
--- a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java
@@ -304,6 +304,9 @@ public class PyTypeHintGenerationUtil {
         collectImportTargetsFromType(pyType, classes, names);
       }
     }
+    if (type instanceof PyInstantiableType && ((PyInstantiableType)type).isDefinition()) {
+      names.add("Type");
+    }
   }
 
   public static final class AnnotationInfo {
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationClassObjectType.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationClassObjectType.py
new file mode 100644
index 000000000000..1f4dfacdd524
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationClassObjectType.py
@@ -0,0 +1,5 @@
+class Foo:
+    pass
+
+
+va<caret>r = Foo
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationClassObjectType_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationClassObjectType_after.py
new file mode 100644
index 000000000000..884c1b9de5f5
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationClassObjectType_after.py
@@ -0,0 +1,8 @@
+from typing import Type
+
+
+class Foo:
+    pass
+
+
+var: [Type[Foo]] = Foo
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
index 044a0b653b6d..4d901949bc03 100644
--- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
+++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
@@ -199,6 +199,10 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase {
     doAnnotationTest();
   }
 
+  public void testAnnotationClassObjectType() {
+    doAnnotationTest();
+  }
+
   private void doAnnotationTest() {
     doTest(LanguageLevel.PYTHON36);
   }