diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java index d903f0d86a76..d4ae63d2a78a 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java @@ -214,20 +214,20 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { private static void insertVariableAnnotation(@NotNull PyTargetExpression target) { final TypeEvalContext context = TypeEvalContext.userInitiated(target.getProject(), target.getContainingFile()); final PyType inferredType = context.getType(target); - PyTypeHintGenerationUtil.checkPep484Compatibility(inferredType); - final String annotationText = PythonDocumentationProvider.getTypeName(inferredType, context); + PyTypeHintGenerationUtil.checkPep484Compatibility(inferredType, context); + final String annotationText = PythonDocumentationProvider.getTypeHint(inferredType, context); final AnnotationInfo info = new AnnotationInfo(annotationText, inferredType); if (isInstanceAttribute(target, context)) { final List classLevelAttrs = findClassLevelDefinitions(target, context); if (classLevelAttrs.isEmpty()) { - PyTypeHintGenerationUtil.insertStandaloneAttributeAnnotation(target, info, true); + PyTypeHintGenerationUtil.insertStandaloneAttributeAnnotation(target, context, info, true); } else { - PyTypeHintGenerationUtil.insertVariableAnnotation(classLevelAttrs.get(0), info, true); + PyTypeHintGenerationUtil.insertVariableAnnotation(classLevelAttrs.get(0), context, info, true); } } else { - PyTypeHintGenerationUtil.insertVariableAnnotation(target, info, true); + PyTypeHintGenerationUtil.insertVariableAnnotation(target, context, info, true); } } @@ -237,15 +237,15 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { if (isInstanceAttribute(target, context)) { final List classLevelAttrs = findClassLevelDefinitions(target, context); if (classLevelAttrs.isEmpty()) { - PyTypeHintGenerationUtil.insertStandaloneAttributeTypeComment(target, info, true); + PyTypeHintGenerationUtil.insertStandaloneAttributeTypeComment(target, context, info, true); } else { // Use existing class level definition (say, assignment of the default value) for annotation - PyTypeHintGenerationUtil.insertVariableTypeComment(classLevelAttrs.get(0), info, true); + PyTypeHintGenerationUtil.insertVariableTypeComment(classLevelAttrs.get(0), context, info, true); } } else { - PyTypeHintGenerationUtil.insertVariableTypeComment(target, info, true); + PyTypeHintGenerationUtil.insertVariableTypeComment(target, context, info, true); } } @@ -285,8 +285,8 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { } else if (target instanceof PyTypedElement) { final PyType singleTargetType = context.getType((PyTypedElement)target); - PyTypeHintGenerationUtil.checkPep484Compatibility(singleTargetType); - final String singleTargetAnnotation = PythonDocumentationProvider.getTypeName(singleTargetType, context); + PyTypeHintGenerationUtil.checkPep484Compatibility(singleTargetType, context); + final String singleTargetAnnotation = PythonDocumentationProvider.getTypeHint(singleTargetType, context); types.add(singleTargetType); typeRanges.add(TextRange.from(builder.length(), singleTargetAnnotation.length())); builder.append(singleTargetAnnotation); diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java index 9b8f0046cc57..bdc08126d67a 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java @@ -48,11 +48,11 @@ public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseInt if (typeComment.getParent() instanceof PyAssignmentStatement && map.size() == 1) { final PyTargetExpression target = ContainerUtil.getFirstItem(map.keySet()); assert target != null; - PyTypeHintGenerationUtil.insertVariableAnnotation(target, new AnnotationInfo(map.get(target)), false); + PyTypeHintGenerationUtil.insertVariableAnnotation(target, null, new AnnotationInfo(map.get(target)), false); } else { for (Map.Entry entry : new ArrayList<>(map.entrySet())) { - PyTypeHintGenerationUtil.insertVariableAnnotation(entry.getKey(), new AnnotationInfo(entry.getValue()), false); + PyTypeHintGenerationUtil.insertVariableAnnotation(entry.getKey(), null, new AnnotationInfo(entry.getValue()), false); } } diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java index 3dc83a8237f4..05b7efba8d53 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java @@ -17,6 +17,7 @@ import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.*; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.ThrowableRunnable; +import com.intellij.util.containers.ContainerUtil; import com.jetbrains.python.PyBundle; import com.jetbrains.python.codeInsight.imports.AddImportHelper; import com.jetbrains.python.codeInsight.imports.AddImportHelper.ImportPriority; @@ -41,7 +42,10 @@ public class PyTypeHintGenerationUtil { private PyTypeHintGenerationUtil() {} - public static void insertStandaloneAttributeTypeComment(@NotNull PyTargetExpression target, AnnotationInfo info, boolean startTemplate) { + public static void insertStandaloneAttributeTypeComment(@NotNull PyTargetExpression target, + @NotNull TypeEvalContext context, + AnnotationInfo info, + boolean startTemplate) { final PyClass pyClass = target.getContainingClass(); if (pyClass == null) { @@ -62,7 +66,7 @@ public class PyTypeHintGenerationUtil { PsiComment insertedComment = as(inserted.getLastChild(), PsiComment.class); if (insertedComment == null) return; - addImportsForTypeAnnotations(info.getTypes(), target.getContainingFile()); + addImportsForTypeAnnotations(info.getTypes(), context, target.getContainingFile()); insertedComment = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(insertedComment); if (startTemplate && insertedComment != null) { @@ -72,6 +76,7 @@ public class PyTypeHintGenerationUtil { } public static void insertStandaloneAttributeAnnotation(@NotNull PyTargetExpression target, + @NotNull TypeEvalContext context, @NotNull AnnotationInfo info, boolean startTemplate) { final LanguageLevel langLevel = LanguageLevel.forElement(target); @@ -94,7 +99,7 @@ public class PyTypeHintGenerationUtil { WriteAction.run(() -> { PyTypeDeclarationStatement inserted = (PyTypeDeclarationStatement)pyClass.getStatementList().addAfter(declaration, anchorBefore); - addImportsForTypeAnnotations(info.getTypes(), target.getContainingFile()); + addImportsForTypeAnnotations(info.getTypes(), context, target.getContainingFile()); inserted = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(inserted); if (startTemplate && inserted != null) { @@ -113,7 +118,10 @@ public class PyTypeHintGenerationUtil { return null; } - public static void insertVariableAnnotation(@NotNull PyTargetExpression target, @NotNull AnnotationInfo info, boolean startTemplate) { + public static void insertVariableAnnotation(@NotNull PyTargetExpression target, + @Nullable TypeEvalContext context, + @NotNull AnnotationInfo info, + boolean startTemplate) { final LanguageLevel langLevel = LanguageLevel.forElement(target); if (langLevel.isOlderThan(LanguageLevel.PYTHON36)) { throw new IllegalArgumentException("Target '" + target.getText() + "' doesn't belong to Python 3.6+ project: " + langLevel); @@ -147,7 +155,9 @@ public class PyTypeHintGenerationUtil { PyAnnotationOwner createdAnnotationOwner = addOrUpdateAnnotatedStatement.compute(); if (createdAnnotationOwner == null) return; - addImportsForTypeAnnotations(info.getTypes(), target.getContainingFile()); + if (context != null) { + addImportsForTypeAnnotations(info.getTypes(), context, target.getContainingFile()); + } createdAnnotationOwner = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(createdAnnotationOwner); if (startTemplate && createdAnnotationOwner != null) { @@ -182,7 +192,10 @@ public class PyTypeHintGenerationUtil { return assignment != null && assignment.getRawTargets().length == 1 && assignment.getLeftHandSideExpression() == target; } - public static void insertVariableTypeComment(@NotNull PyTargetExpression target, @NotNull AnnotationInfo info, boolean startTemplate) { + public static void insertVariableTypeComment(@NotNull PyTargetExpression target, + TypeEvalContext context, + @NotNull AnnotationInfo info, + boolean startTemplate) { if (!FileModificationService.getInstance().preparePsiElementForWrite(target)) return; final String typeCommentText = " " + TYPE_COMMENT_PREFIX + info.getAnnotationText(); @@ -231,7 +244,7 @@ public class PyTypeHintGenerationUtil { PsiComment insertedComment = target.getTypeComment(); if (insertedComment == null) return; - addImportsForTypeAnnotations(info.getTypes(), target.getContainingFile()); + addImportsForTypeAnnotations(info.getTypes(), context, target.getContainingFile()); insertedComment = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(insertedComment); if (startTemplate && insertedComment != null) { @@ -264,12 +277,14 @@ public class PyTypeHintGenerationUtil { } } - public static void addImportsForTypeAnnotations(@NotNull List types, @NotNull PsiFile file) { + private static void addImportsForTypeAnnotations(@NotNull List types, + @NotNull TypeEvalContext context, + @NotNull PsiFile file) { final Set classes = new HashSet<>(); final Set namesFromTyping = new HashSet<>(); for (PyType type : types) { - collectImportTargetsFromType(type, classes, namesFromTyping); + collectImportTargetsFromType(type, context, classes, namesFromTyping); } final boolean builtinTyping = LanguageLevel.forElement(file).isAtLeast(LanguageLevel.PYTHON35); @@ -283,7 +298,10 @@ public class PyTypeHintGenerationUtil { } } - private static void collectImportTargetsFromType(@Nullable PyType type, @NotNull Set classes, @NotNull Set names) { + private static void collectImportTargetsFromType(@Nullable PyType type, + @NotNull TypeEvalContext context, + @NotNull Set classes, + @NotNull Set names) { if (type == null) { names.add("Any"); } @@ -292,7 +310,7 @@ public class PyTypeHintGenerationUtil { final boolean isOptional = members.size() == 2 && members.contains(PyNoneType.INSTANCE); names.add(isOptional ? "Optional" : "Union"); for (PyType pyType : members) { - collectImportTargetsFromType(pyType, classes, names); + collectImportTargetsFromType(pyType, context, classes, names); } } else if (type instanceof PyCollectionType) { @@ -310,40 +328,54 @@ public class PyTypeHintGenerationUtil { names.add("Tuple"); } for (PyType pyType : ((PyCollectionType)type).getElementTypes()) { - collectImportTargetsFromType(pyType, classes, names); + collectImportTargetsFromType(pyType, context, classes, names); } } else if (type instanceof PyClassType) { classes.add(((PyClassType)type).getPyClass()); } + else if (type instanceof PyCallableType) { + names.add("Callable"); + final PyCallableType callableType = (PyCallableType)type; + for (PyCallableParameter parameter : ContainerUtil.notNullize(callableType.getParameters(context))) { + collectImportTargetsFromType(parameter.getType(context), context, classes, names); + } + collectImportTargetsFromType(callableType.getReturnType(context), context, classes, names); + } if (type instanceof PyInstantiableType && ((PyInstantiableType)type).isDefinition()) { names.add("Type"); } } - public static void checkPep484Compatibility(@Nullable PyType type) { + public static void checkPep484Compatibility(@Nullable PyType type, @NotNull TypeEvalContext context) { if (type == null || type instanceof PyNoneType || - type instanceof PyCallableTypeImpl || - type instanceof PyGenericType || - type instanceof PyFunctionType) { + type instanceof PyGenericType) { return; } else if (type instanceof PyUnionType) { for (PyType memberType : ((PyUnionType)type).getMembers()) { - checkPep484Compatibility(memberType); + checkPep484Compatibility(memberType, context); } } else if (type instanceof PyCollectionType) { for (PyType typeParam : ((PyCollectionType)type).getElementTypes()) { - checkPep484Compatibility(typeParam); + checkPep484Compatibility(typeParam, context); } } else if (type instanceof PyClassType) { // In this order since PyCollectionTypeImpl implements PyClassType } + else if (type instanceof PyCallableType) { + final PyCallableType callableType = (PyCallableType)type; + for (PyCallableParameter parameter : ContainerUtil.notNullize(callableType.getParameters(context))) { + checkPep484Compatibility(parameter.getType(context), context); + } + checkPep484Compatibility(callableType.getReturnType(context), context); + } else { - throw new Pep484IncompatibleTypeException(PyBundle.message("INTN.add.type.hint.for.variable.PEP484.incompatible.type", type.getName())); + throw new Pep484IncompatibleTypeException( + PyBundle.message("INTN.add.type.hint.for.variable.PEP484.incompatible.type", type.getName())); } } diff --git a/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java b/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java index 50132c086242..8b1658237aef 100644 --- a/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java +++ b/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java @@ -51,6 +51,13 @@ public class PyTypeModelBuilder { accept(visitor); return visitor.getDescription(); } + + @NotNull + public String asPep484TypeHint() { + final TypeToStringVisitor visitor = new TypeToPep484TypeHintVisitor(); + accept(visitor); + return visitor.getString(); + } } static class OneOf extends TypeModel { @@ -380,6 +387,40 @@ public class PyTypeModelBuilder { } } + private static class TypeToPep484TypeHintVisitor extends TypeToStringVisitor { + @Override + protected boolean maxDepthExceeded() { + return false; + } + + @Override + public void function(FunctionType function) { + add("Callable["); + final Collection parameters = function.parameters; + if (parameters != null) { + add("["); + processList(parameters); + add("]"); + } + else { + add("..."); + } + add(", "); + function.returnType.accept(this); + add("]"); + } + + @Override + public void param(ParamType param) { + if (param.type != null) { + param.type.accept(this); + } + else { + add("Any"); + } + } + } + private static class TypeToBodyWithLinksVisitor extends TypeNameVisitor { private ChainIterable myBody; private PsiElement myAnchor; @@ -429,7 +470,7 @@ public class PyTypeModelBuilder { @Override public void oneOf(OneOf oneOf) { myDepth++; - if (myDepth > MAX_DEPTH) { + if (maxDepthExceeded()) { add("..."); return; } @@ -439,7 +480,7 @@ public class PyTypeModelBuilder { myDepth--; } - private void processList(@NotNull Collection list) { + protected void processList(@NotNull Collection list) { boolean first = true; for (TypeModel t : list) { if (!first) { @@ -458,7 +499,7 @@ public class PyTypeModelBuilder { @Override public void collectionOf(CollectionOf collectionOf) { myDepth++; - if (myDepth > MAX_DEPTH) { + if (maxDepthExceeded()) { add("..."); return; } @@ -481,7 +522,7 @@ public class PyTypeModelBuilder { @Override public void function(FunctionType function) { myDepth++; - if (myDepth > MAX_DEPTH) { + if (maxDepthExceeded()) { add("..."); return; } @@ -498,10 +539,14 @@ public class PyTypeModelBuilder { myDepth--; } + protected boolean maxDepthExceeded() { + return myDepth > MAX_DEPTH; + } + @Override public void param(ParamType param) { myDepth++; - if (myDepth > MAX_DEPTH) { + if (maxDepthExceeded()) { add("..."); return; } diff --git a/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java b/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java index 1bcc692162a4..d5f473d7319c 100644 --- a/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java +++ b/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java @@ -246,6 +246,14 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i return buildTypeModel(type, context).asString(); } + /** + * Returns the provided type in PEP 484 compliant format. + */ + @NotNull + public static String getTypeHint(@Nullable PyType type, @NotNull TypeEvalContext context) { + return buildTypeModel(type, context).asPep484TypeHint(); + } + /** * @param type type which description will be calculated. * Description is the same as {@link PythonDocumentationProvider#getTypeDescription(PyType, TypeEvalContext)} gives but diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType_after.py index 74fb17c6a316..6423a823779e 100644 --- a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType_after.py +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType_after.py @@ -5,4 +5,4 @@ def func(x): pass -var: Callable[[Any], None] = func +var: [Callable[[Any], None]] = func \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationNotPossibleForStructuralTypeInCallable.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationNotPossibleForStructuralTypeInCallable.py new file mode 100644 index 000000000000..f4cd03696190 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationNotPossibleForStructuralTypeInCallable.py @@ -0,0 +1,6 @@ +def func(x): + x.foo() + return x + + +var = func \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationNotPossibleForStructuralTypeInCallable_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationNotPossibleForStructuralTypeInCallable_after.py new file mode 100644 index 000000000000..55b25b8fa72d --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationNotPossibleForStructuralTypeInCallable_after.py @@ -0,0 +1,6 @@ +def func(x): + x.foo() + return x + + +var = func \ 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 a5ee7a693a97..f8f32eefab5b 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java @@ -223,6 +223,10 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { doAnnotationTest(); } + public void testAnnotationNotPossibleForStructuralTypeInCallable() { + doAnnotationTest(); + } + public void testAnnotationCallableType() { doAnnotationTest(); }