diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java index 05248c5ec2fd..288ed125efec 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java @@ -9,9 +9,11 @@ import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.IncorrectOperationException; import com.jetbrains.python.PyBundle; import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider; +import com.jetbrains.python.documentation.PythonDocumentationProvider; import com.jetbrains.python.documentation.doctest.PyDocstringFile; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.resolve.PyResolveContext; +import com.jetbrains.python.psi.types.PyType; import com.jetbrains.python.psi.types.TypeEvalContext; import one.util.streamex.StreamEx; import org.jetbrains.annotations.Nls; @@ -74,7 +76,7 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { final List targets = findSuitableTargetsUnderCaret(project, editor, file); assert targets.size() == 1; final PyTargetExpression annotationTarget = targets.get(0); - if (LanguageLevel.forElement(annotationTarget).isAtLeast(LanguageLevel.PYTHON36)) { + if (preferSyntacticAnnotation(annotationTarget)) { insertVariableAnnotation(annotationTarget); } else { @@ -82,8 +84,15 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { } } - private void insertVariableAnnotation(@NotNull PyTargetExpression target) { + private static boolean preferSyntacticAnnotation(@NotNull PyTargetExpression annotationTarget) { + return LanguageLevel.forElement(annotationTarget).isAtLeast(LanguageLevel.PYTHON36); + } + private static void insertVariableAnnotation(@NotNull PyTargetExpression target) { + final TypeEvalContext context = TypeEvalContext.userInitiated(target.getProject(), target.getContainingFile()); + final PyType inferredType = context.getType(target); + final String annotationText = PythonDocumentationProvider.getTypeName(inferredType, context); + PyTypeHintGenerationUtil.insertVariableAnnotation(target, annotationText); } private void insertVariableTypeComment(@NotNull PyTargetExpression target) { diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java index 106309ca4ee8..134a3a1aa955 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java @@ -35,7 +35,6 @@ import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import java.util.ArrayList; -import java.util.List; import java.util.Map; public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseIntentionAction { @@ -48,24 +47,11 @@ public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseInt if (typeComment.getParent() instanceof PyAssignmentStatement && map.size() == 1) { final PyTargetExpression target = ContainerUtil.getFirstItem(map.keySet()); assert target != null; - PyUtil.updateDocumentUnblockedAndCommitted(target, document -> { - document.insertString(target.getTextRange().getEndOffset(), ": " + map.get(target)); - }); + PyTypeHintGenerationUtil.insertVariableAnnotation(target, map.get(target)); } else { - final PyStatement statement = PsiTreeUtil.getParentOfType(typeComment, PyStatement.class); - assert statement != null; - - final PyElementGenerator generator = PyElementGenerator.getInstance(project); - final List> entries = new ArrayList<>(map.entrySet()); - - for (Map.Entry entry : entries) { - final PyTargetExpression target = entry.getKey(); - final String annotation = entry.getValue(); - final PyTypeDeclarationStatement declaration = generator.createFromText(LanguageLevel.PYTHON36, - PyTypeDeclarationStatement.class, - target.getText() + ": " + annotation); - statement.getParent().addBefore(declaration, statement); + for (Map.Entry entry : new ArrayList<>(map.entrySet())) { + PyTypeHintGenerationUtil.insertVariableAnnotation(entry.getKey(), entry.getValue()); } } diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java new file mode 100644 index 000000000000..9c0278101c6a --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java @@ -0,0 +1,41 @@ +// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. +package com.jetbrains.python.codeInsight.intentions; + +import com.intellij.psi.util.PsiTreeUtil; +import com.jetbrains.python.psi.*; +import org.jetbrains.annotations.NotNull; + +import static com.jetbrains.python.psi.PyUtil.as; + +/** + * @author Mikhail Golubev + */ +public class PyTypeHintGenerationUtil { + private PyTypeHintGenerationUtil() {} + + public static void insertVariableAnnotation(@NotNull PyTargetExpression target, @NotNull String annotation) { + 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); + } + + if (canUseInlineAnnotation(target)) { + PyUtil.updateDocumentUnblockedAndCommitted(target, document -> { + document.insertString(target.getTextRange().getEndOffset(), ": " + annotation); + }); + } + else { + final PyElementGenerator generator = PyElementGenerator.getInstance(target.getProject()); + final String declarationText = target.getText() + ": " + annotation; + final PyTypeDeclarationStatement declaration = generator.createFromText(langLevel, PyTypeDeclarationStatement.class, declarationText); + final PyStatement statement = PsiTreeUtil.getParentOfType(target, PyStatement.class); + assert statement != null; + statement.getParent().addBefore(declaration, statement); + } + } + + private static boolean canUseInlineAnnotation(@NotNull PyTargetExpression target) { + final PyAssignmentStatement assignment = as(target.getParent(), PyAssignmentStatement.class); + return assignment != null && assignment.getRawTargets().length == 1 && assignment.getLeftHandSideExpression() == target; + } +} diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalChainedAssignmentTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalChainedAssignmentTarget.py new file mode 100644 index 000000000000..e1b3a1ccf2fd --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalChainedAssignmentTarget.py @@ -0,0 +1,3 @@ +def func(): + var = another = 'spam' + var \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalChainedAssignmentTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalChainedAssignmentTarget_after.py new file mode 100644 index 000000000000..bb3bd7b8af67 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalChainedAssignmentTarget_after.py @@ -0,0 +1,4 @@ +def func(): + var: str + var = another = 'spam' + var \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalForTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalForTarget.py new file mode 100644 index 000000000000..538221a6b634 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalForTarget.py @@ -0,0 +1,3 @@ +def func(): + for var in 'spam': + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalForTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalForTarget_after.py new file mode 100644 index 000000000000..b20adfd248ad --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalForTarget_after.py @@ -0,0 +1,4 @@ +def func(): + var: str + for var in 'spam': + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTarget.py similarity index 100% rename from python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget.py rename to python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTarget.py diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTargetInParentheses.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTargetInParentheses.py new file mode 100644 index 000000000000..820891b8cb69 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTargetInParentheses.py @@ -0,0 +1,3 @@ +def func(): + (var) = 'spam' + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTargetInParentheses_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTargetInParentheses_after.py new file mode 100644 index 000000000000..3cf034e02654 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTargetInParentheses_after.py @@ -0,0 +1,4 @@ +def func(): + var: Any + (var) = 'spam' + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTarget_after.py similarity index 100% rename from python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget_after.py rename to python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalSimpleAssignmentTarget_after.py diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalUnpackedAssignmentTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalUnpackedAssignmentTarget.py new file mode 100644 index 000000000000..05cb637ffc2c --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalUnpackedAssignmentTarget.py @@ -0,0 +1,3 @@ +def func(): + var, _ = 'spam', 42 + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalUnpackedAssignmentTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalUnpackedAssignmentTarget_after.py new file mode 100644 index 000000000000..a4f02609aefc --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalUnpackedAssignmentTarget_after.py @@ -0,0 +1,4 @@ +def func(): + var: str + var, _ = 'spam', 42 + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalWithTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalWithTarget.py new file mode 100644 index 000000000000..06db7a3f427f --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalWithTarget.py @@ -0,0 +1,3 @@ +def func(): + with open('file.txt') as var: + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalWithTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalWithTarget_after.py new file mode 100644 index 000000000000..d351d8014aa9 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalWithTarget_after.py @@ -0,0 +1,4 @@ +def func(): + var: TextIO + with open('file.txt') as var: + var diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java index e033b5c1dd39..ff6c59ee88c2 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java @@ -9,7 +9,27 @@ import org.jetbrains.annotations.NotNull; * @author Mikhail Golubev */ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { - public void testAnnotationLocalAssignmentTarget() { + public void testAnnotationLocalSimpleAssignmentTarget() { + doTestAnnotation(); + } + + public void testAnnotationLocalSimpleAssignmentTargetInParentheses() { + doTestAnnotation(); + } + + public void testAnnotationLocalUnpackedAssignmentTarget() { + doTestAnnotation(); + } + + public void testAnnotationLocalChainedAssignmentTarget() { + doTestAnnotation(); + } + + public void testAnnotationLocalForTarget() { + doTestAnnotation(); + } + + public void testAnnotationLocalWithTarget() { doTestAnnotation(); }