diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java index 288ed125efec..53d503174ff2 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java @@ -95,7 +95,10 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { PyTypeHintGenerationUtil.insertVariableAnnotation(target, annotationText); } - private void insertVariableTypeComment(@NotNull PyTargetExpression target) { - + private static void insertVariableTypeComment(@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.insertVariableTypeComment(target, annotationText); } } diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java index 9c0278101c6a..73ddc520c011 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java @@ -38,4 +38,23 @@ public class PyTypeHintGenerationUtil { final PyAssignmentStatement assignment = as(target.getParent(), PyAssignmentStatement.class); return assignment != null && assignment.getRawTargets().length == 1 && assignment.getLeftHandSideExpression() == target; } + + public static void insertVariableTypeComment(@NotNull PyTargetExpression target, @NotNull String annotation) { + final PyStatement statement = PsiTreeUtil.getParentOfType(target, PyStatement.class); + final String typeCommentText = " # type: " + annotation; + if (statement instanceof PyAssignmentStatement) { + PyUtil.updateDocumentUnblockedAndCommitted(target, document -> { + document.insertString(statement.getTextRange().getEndOffset(), typeCommentText); + }); + } + else if (statement instanceof PyWithStatement || statement instanceof PyForStatement) { + PyUtil.updateDocumentUnblockedAndCommitted(target, document -> { + final PyStatementListContainer container = statement instanceof PyForStatement ? + ((PyForStatement)statement).getForPart() : + (PyWithStatement)statement; + final int endOffset = PyUtil.getHeaderEndAnchor(container).getTextRange().getEndOffset(); + document.insertString(endOffset, typeCommentText); + }); + } + } } diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index e0ccc25c53e0..a15dd29953ce 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/python/src/com/jetbrains/python/psi/PyUtil.java @@ -865,8 +865,13 @@ public class PyUtil { */ @Nullable public static PsiComment getCommentOnHeaderLine(@NotNull PyStatementListContainer container) { + return as(getHeaderEndAnchor(container), PsiComment.class); + } + + @NotNull + public static PsiElement getHeaderEndAnchor(@NotNull PyStatementListContainer container) { final PyStatementList statementList = container.getStatementList(); - return as(PyPsiUtils.getPrevNonWhitespaceSibling(statementList), PsiComment.class); + return ObjectUtils.notNull(PyPsiUtils.getPrevNonWhitespaceSibling(statementList)); } public static boolean isPy2ReservedWord(@NotNull PyReferenceExpression node) { diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalChainedAssignmentTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalChainedAssignmentTarget.py new file mode 100644 index 000000000000..e1b3a1ccf2fd --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalChainedAssignmentTarget.py @@ -0,0 +1,3 @@ +def func(): + var = another = 'spam' + var \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalChainedAssignmentTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalChainedAssignmentTarget_after.py new file mode 100644 index 000000000000..831ed3665c3a --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalChainedAssignmentTarget_after.py @@ -0,0 +1,3 @@ +def func(): + var = another = 'spam' # type: str + var \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalForTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalForTarget.py new file mode 100644 index 000000000000..538221a6b634 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalForTarget.py @@ -0,0 +1,3 @@ +def func(): + for var in 'spam': + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalForTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalForTarget_after.py new file mode 100644 index 000000000000..93d47f20b721 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalForTarget_after.py @@ -0,0 +1,3 @@ +def func(): + for var in 'spam': # type: str + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTarget.py new file mode 100644 index 000000000000..bfdb8566d32c --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTarget.py @@ -0,0 +1,3 @@ +def func(): + var = 'spam' + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTargetInParentheses.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTargetInParentheses.py new file mode 100644 index 000000000000..820891b8cb69 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTargetInParentheses.py @@ -0,0 +1,3 @@ +def func(): + (var) = 'spam' + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTargetInParentheses_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTargetInParentheses_after.py new file mode 100644 index 000000000000..07928c0317bb --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTargetInParentheses_after.py @@ -0,0 +1,3 @@ +def func(): + (var) = 'spam' # type: Any + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTarget_after.py new file mode 100644 index 000000000000..765e9d1db6ea --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalSimpleAssignmentTarget_after.py @@ -0,0 +1,3 @@ +def func(): + var = 'spam' # type: str + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalUnpackedAssignmentTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalUnpackedAssignmentTarget.py new file mode 100644 index 000000000000..05cb637ffc2c --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalUnpackedAssignmentTarget.py @@ -0,0 +1,3 @@ +def func(): + var, _ = 'spam', 42 + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalUnpackedAssignmentTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalUnpackedAssignmentTarget_after.py new file mode 100644 index 000000000000..2c573ff21f6a --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalUnpackedAssignmentTarget_after.py @@ -0,0 +1,3 @@ +def func(): + var, _ = 'spam', 42 # type: (str, int) + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalWithTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalWithTarget.py new file mode 100644 index 000000000000..06db7a3f427f --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalWithTarget.py @@ -0,0 +1,3 @@ +def func(): + with open('file.txt') as var: + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalWithTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalWithTarget_after.py new file mode 100644 index 000000000000..fdd8ca797bbc --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentLocalWithTarget_after.py @@ -0,0 +1,3 @@ +def func(): + with open('file.txt') as var: # type: BinaryIO + var diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java index ff6c59ee88c2..e3bbbfeaac3d 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java @@ -33,6 +33,30 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { doTestAnnotation(); } + public void testTypeCommentLocalSimpleAssignmentTarget() { + doTestTypeComment(); + } + + public void testTypeCommentLocalSimpleAssignmentTargetInParentheses() { + doTestTypeComment(); + } + + public void testTypeCommentLocalUnpackedAssignmentTarget() { + doTestTypeComment(); + } + + public void testTypeCommentLocalChainedAssignmentTarget() { + doTestTypeComment(); + } + + public void testTypeCommentLocalForTarget() { + doTestTypeComment(); + } + + public void testTypeCommentLocalWithTarget() { + doTestTypeComment(); + } + private void doTestAnnotation() { doTest(LanguageLevel.PYTHON36); }