PY-20801 Handle "with" statements and "for" loops by declaring type beforehand

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java
index 8fb2e15d81274fb2e15a83361b58b99de9754773..e8e2d7fc8cde3b68f7f17b50749a4cfc7d2db6e8 100644 (file)
--- 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 (file)
index 0000000..d052c45
--- /dev/null
+++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking.py
@@ -0,0 +1,2 @@
+for x, y in undefined():  # ty<caret>pe: 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 (file)
index 0000000..86b64b1
--- /dev/null
+++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleForLoop.py
@@ -0,0 +1,2 @@
+for x in undefined():  # ty<caret>pe: 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 (file)
index 0000000..c3e537c
--- /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 (file)
index 0000000..a6e56ac
--- /dev/null
+++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement.py
@@ -0,0 +1,2 @@
+with undefined() as x:  # ty<caret>pe: int
+    pass

diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/simpleWithStatement_after.py
new file mode 100644 (file)
index 0000000..1ee5602
--- /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 (file)
index 0000000..7428426
--- /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 (file)
index 0000000..4c94497
--- /dev/null
+++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking.py
@@ -0,0 +1,2 @@
+with undefined() as (x, y):  # ty<caret>pe: 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 e912da672ef049fd3813780c53bed4809596358f..2bac1e73be14a8cf2b553e902089a3cefff15a53 100644 (file)
--- 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();
+  }
 }