From 5e772eb11bd7e9a477a5f58646ac286e9d194295 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Thu, 20 Oct 2016 18:37:16 +0300 Subject: [PATCH] PY-20801 Handle "with" statements and "for" loops by declaring type beforehand --- ...onvertTypeCommentToVariableAnnotation.java | 73 ++++++++++++++++--- .../forLoopWithUnpacking.py | 2 + .../simpleForLoop.py | 2 + .../simpleForLoop_after.py | 3 + .../simpleWithStatement.py | 2 + .../simpleWithStatement_after.py | 3 + .../withStatementWithMultipleWithItems.py | 2 + .../withStatementWithUnpacking.py | 2 + ...mentToVariableAnnotationIntentionTest.java | 20 +++++ 9 files changed, 98 insertions(+), 11 deletions(-) create mode 100644 python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking.py create mode 100644 python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop.py create mode 100644 python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop_after.py create mode 100644 python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement.py create mode 100644 python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement_after.py create mode 100644 python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithMultipleWithItems.py create mode 100644 python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking.py diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java index 8fb2e15d8127..e8e2d7fc8cde 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java @@ -32,22 +32,40 @@ import org.jetbrains.annotations.Nls; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; +import static com.jetbrains.python.psi.PyUtil.as; + public class PyConvertTypeCommentToVariableAnnotation extends PyBaseIntentionAction { @Override public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { final PsiComment comment = findCommentUnderCaret(editor, file); if (comment != null) { - final PyTargetExpression target = findTypeCommentTarget(comment); - if (target != null) { - final String annotation = PyTypingTypeProvider.getTypeCommentValue(comment.getText()); - final PsiElement prev = PyPsiUtils.getPrevNonWhitespaceSibling(comment); - final int commentStart = prev != null ? prev.getTextRange().getEndOffset() : comment.getTextRange().getStartOffset(); - final int commentEnd = comment.getTextRange().getEndOffset(); + final String annotation = PyTypingTypeProvider.getTypeCommentValue(comment.getText()); + final PyTargetExpression assignmentTarget = findAssignmentTypeCommentTarget(comment); + if (assignmentTarget != null) { + comment.delete(); final Document document = editor.getDocument(); runWithDocumentReleasedAndCommitted(project, document, () -> { - document.deleteString(commentStart, commentEnd); - document.insertString(target.getTextRange().getEndOffset(), ": " + annotation); + document.insertString(assignmentTarget.getTextRange().getEndOffset(), ": " + annotation); }); + return; + } + final PyTargetExpression compoundTarget; + final PyTargetExpression forTarget = findForLoopTypeCommentTarget(comment); + if (forTarget != null) { + compoundTarget = forTarget; + } + else { + compoundTarget = findWithStatementTypeCommentTarget(comment); + } + if (compoundTarget != null) { + comment.delete(); + final PyElementGenerator generator = PyElementGenerator.getInstance(project); + final PyTypeDeclarationStatement declaration = generator.createFromText(LanguageLevel.PYTHON36, + PyTypeDeclarationStatement.class, + compoundTarget.getText() + ": " + annotation); + final PyStatement containingStatement = PsiTreeUtil.getParentOfType(compoundTarget, PyStatement.class); + assert containingStatement != null; + containingStatement.getParent().addBefore(declaration, containingStatement); } } } @@ -80,13 +98,15 @@ public class PyConvertTypeCommentToVariableAnnotation extends PyBaseIntentionAct return PsiTreeUtil.getParentOfType(element, PsiComment.class, false); } - private boolean isSuitableTypeComment(@NotNull PsiComment comment) { + private static boolean isSuitableTypeComment(@NotNull PsiComment comment) { final String annotation = PyTypingTypeProvider.getTypeCommentValue(comment.getText()); - return annotation != null && findTypeCommentTarget(comment) != null; + return annotation != null && (findAssignmentTypeCommentTarget(comment) != null || + findForLoopTypeCommentTarget(comment) != null || + findWithStatementTypeCommentTarget(comment) != null); } @Nullable - private PyTargetExpression findTypeCommentTarget(@NotNull PsiComment comment) { + private static PyTargetExpression findAssignmentTypeCommentTarget(@NotNull PsiComment comment) { final PsiElement parent = comment.getParent(); if (parent instanceof PyAssignmentStatement) { final PyAssignmentStatement assignment = (PyAssignmentStatement)parent; @@ -101,6 +121,37 @@ public class PyConvertTypeCommentToVariableAnnotation extends PyBaseIntentionAct return null; } + @Nullable + private static PyTargetExpression findForLoopTypeCommentTarget(@NotNull PsiComment comment) { + final PsiElement parent = comment.getParent(); + if (parent instanceof PyForPart) { + final PyForPart forPart = (PyForPart)parent; + final PyTargetExpression target = as(forPart.getTarget(), PyTargetExpression.class); + if (target != null && target.getTypeComment() == comment) { + return target; + } + } + return null; + } + + @Nullable + private static PyTargetExpression findWithStatementTypeCommentTarget(@NotNull PsiComment comment) { + final PsiElement parent = comment.getParent(); + if (parent instanceof PyWithStatement) { + final PyWithStatement withStatement = (PyWithStatement)parent; + final PyWithItem[] withItems = withStatement.getWithItems(); + if (withItems.length == 1) { + final PyTargetExpression target = as(withItems[0].getTarget(), PyTargetExpression.class); + if (target != null && target.getTypeComment() == comment) { + return target; + } + } + } + return null; + } + + + public static void runWithDocumentReleasedAndCommitted(@NotNull Project project, @NotNull Document document, @NotNull Runnable runnable) { final PsiDocumentManager manager = PsiDocumentManager.getInstance(project); manager.doPostponedOperationsAndUnblockDocument(document); diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking.py new file mode 100644 index 000000000000..d052c452a3b6 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking.py @@ -0,0 +1,2 @@ +for x, y in undefined(): # type: int, str + pass \ No newline at end of file diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop.py new file mode 100644 index 000000000000..86b64b14c2c3 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop.py @@ -0,0 +1,2 @@ +for x in undefined(): # type: int + pass \ No newline at end of file diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop_after.py new file mode 100644 index 000000000000..c3e537c8aabb --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop_after.py @@ -0,0 +1,3 @@ +x: int +for x in undefined(): + pass \ No newline at end of file diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement.py new file mode 100644 index 000000000000..a6e56ac3305e --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement.py @@ -0,0 +1,2 @@ +with undefined() as x: # type: int + pass diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement_after.py new file mode 100644 index 000000000000..1ee5602a74b4 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement_after.py @@ -0,0 +1,3 @@ +x: int +with undefined() as x: + pass diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithMultipleWithItems.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithMultipleWithItems.py new file mode 100644 index 000000000000..742842677a97 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithMultipleWithItems.py @@ -0,0 +1,2 @@ +with undefined() as x, undefined() as y: # type: int + pass diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking.py new file mode 100644 index 000000000000..4c9449793a23 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking.py @@ -0,0 +1,2 @@ +with undefined() as (x, y): # type: int, str + pass \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest.java index e912da672ef0..2bac1e73be14 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest.java @@ -50,4 +50,24 @@ public class PyConvertTypeCommentToVariableAnnotationIntentionTest extends PyInt public void testMultilineAssignment() { doPositiveTest(); } + + public void testSimpleForLoop() { + doPositiveTest(); + } + + public void testSimpleWithStatement() { + doPositiveTest(); + } + + public void testForLoopWithUnpacking() { + doNegativeTest(); + } + + public void testWithStatementWithUnpacking() { + doNegativeTest(); + } + + public void testWithStatementWithMultipleWithItems() { + doNegativeTest(); + } }