diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java index 992219572ede..9ec23b4b8fe1 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java @@ -19,6 +19,7 @@ import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.ThrowableRunnable; import com.jetbrains.python.codeInsight.imports.AddImportHelper; import com.jetbrains.python.codeInsight.imports.AddImportHelper.ImportPriority; +import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyPsiUtils; import com.jetbrains.python.psi.types.*; @@ -293,17 +294,24 @@ public class PyTypeHintGenerationUtil { collectImportTargetsFromType(pyType, classes, names); } } - else if (type instanceof PyClassType) { - classes.add(((PyClassType)type).getPyClass()); - } else if (type instanceof PyCollectionType) { if (type instanceof PyCollectionTypeImpl) { - classes.add(((PyCollectionTypeImpl)type).getPyClass()); + final PyClass pyClass = ((PyCollectionTypeImpl)type).getPyClass(); + final String typingCollectionName = PyTypingTypeProvider.TYPING_COLLECTION_CLASSES.get(pyClass.getQualifiedName()); + if (typingCollectionName != null && type.isBuiltin()) { + names.add(typingCollectionName); + } + else { + classes.add(pyClass); + } } for (PyType pyType : ((PyCollectionType)type).getElementTypes()) { collectImportTargetsFromType(pyType, classes, names); } } + else if (type instanceof PyClassType) { + classes.add(((PyClassType)type).getPyClass()); + } if (type instanceof PyInstantiableType && ((PyInstantiableType)type).isDefinition()) { names.add("Type"); } diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationListType.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationListType.py new file mode 100644 index 000000000000..bede741880e4 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationListType.py @@ -0,0 +1,3 @@ +def func(): + var = [1, 2, 3] + var \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationListType_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationListType_after.py new file mode 100644 index 000000000000..f29ab3206d04 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationListType_after.py @@ -0,0 +1,6 @@ +from typing import List + + +def func(): + var: [List[int]] = [1, 2, 3] + var \ 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 4d901949bc03..f5c649d8ab02 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java @@ -203,6 +203,10 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { doAnnotationTest(); } + public void testAnnotationListType() { + doAnnotationTest(); + } + private void doAnnotationTest() { doTest(LanguageLevel.PYTHON36); }