PY-20801 Handle cases with nested target expressions and unpacking
authorMikhail Golubev <mikhail.golubev@jetbrains.com>
Thu, 20 Oct 2016 21:25:31 +0000 (00:25 +0300)
committerMikhail Golubev <mikhail.golubev@jetbrains.com>
Mon, 24 Oct 2016 21:03:49 +0000 (00:03 +0300)
14 files changed:
python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/after.py.template
python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/before.py.template
python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/description.html
python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithComplexUnpacking.py [new file with mode: 0644]
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithComplexUnpacking_after.py [new file with mode: 0644]
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithUnpacking_after.py [new file with mode: 0644]
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithComplexUnpacking.py [new file with mode: 0644]
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithComplexUnpacking_after.py [new file with mode: 0644]
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking_after.py [new file with mode: 0644]
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithComplexUnpacking.py [new file with mode: 0644]
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithComplexUnpacking_after.py [new file with mode: 0644]
python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking_after.py [new file with mode: 0644]
python/testSrc/com/jetbrains/python/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest.java

index 89023c6848adccceed370537713fefa304b3721a..680d63c083938becbe93eff3c03aaea05c2eb76d 100644 (file)
@@ -1 +1,13 @@
-x: Optional[int] = undefined()
\ No newline at end of file
+x: List[Item] = locals()['foo']
+line: int
+col: int
+line, col = hero.pos
+
+key: str
+value: int
+for key, value in d.items():
+    ...
+
+f: io.FIle
+with open('bar.txt') as f:
+    ...
\ No newline at end of file
index e28b15f6074be2a8b31118c321de01103c55d563..4af9ffca879bc0396513cfa2478756f27bf61751 100644 (file)
@@ -1 +1,8 @@
-x = undefined()  # type: Optional[int]
\ No newline at end of file
+x = locals()['foo']  # type: List[Item]
+line, col = hero.pos  # type: int, int
+
+for key, value in d.items():  # type: str, int
+    ...
+
+with open() as f:  # type: io.FIle
+    ...
index 2d034cc57781f6cbe490ca6893f752642efa4ed6..6d235944ba5e95f4136404027e0d629a03a7e811 100644 (file)
@@ -1,7 +1,7 @@
 <html>
 <body>
 <span>
-  This intention convert PEP-484 type comments to Python 3.6 variable annotations
+  This intention convert PEP-484 type comments to Python 3.6 variable annotations.
 </span>
 </body>
 </html>
\ No newline at end of file
index e8e2d7fc8cde3b68f7f17b50749a4cfc7d2db6e8..9bb6485bc42517004e0efc034a4ecf67e0f888c5 100644 (file)
  */
 package com.jetbrains.python.codeInsight.intentions;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Ref;
 import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
+import com.intellij.util.containers.ContainerUtil;
+import com.intellij.util.containers.hash.LinkedHashMap;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.codeInsight.PyTypingTypeProvider;
 import com.jetbrains.python.psi.*;
@@ -32,40 +37,47 @@ import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.*;
+
 import static com.jetbrains.python.psi.PyUtil.as;
+import static com.jetbrains.python.psi.PyUtil.rehighlightOpenEditors;
 
 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 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.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);
+    final PsiComment typeComment = findCommentUnderCaret(editor, file);
+    if (typeComment != null) {
+      final Map<PyTargetExpression, String> map = mapTargetsToAnnotations(typeComment);
+      if (map != null) {
+        if (typeComment.getParent() instanceof PyAssignmentStatement && map.size() == 1) {
+          final Document document = editor.getDocument();
+          runWithDocumentReleasedAndCommitted(project, document, () -> {
+            final PyTargetExpression target = ContainerUtil.getFirstItem(map.keySet());
+            assert target != null;
+            document.insertString(target.getTextRange().getEndOffset(), ": " + map.get(target));
+          });
+        }
+        else {
+          final PyStatement statement = PsiTreeUtil.getParentOfType(typeComment, PyStatement.class);
+          assert statement != null;
+
+          final PyElementGenerator generator = PyElementGenerator.getInstance(project);
+          final List<Map.Entry<PyTargetExpression, String>> entries = new ArrayList<>(map.entrySet());
+          Collections.reverse(entries);
+
+          PsiElement anchor = statement;
+          for (Map.Entry<PyTargetExpression, String> entry : entries) {
+            final PyTargetExpression target = entry.getKey();
+            final String annotation = entry.getValue();
+            final PyTypeDeclarationStatement declaration = generator.createFromText(LanguageLevel.PYTHON36,
+                                                                                    PyTypeDeclarationStatement.class,
+                                                                                    target.getText() + ": " + annotation);
+            anchor = statement.getParent().addBefore(declaration, anchor);
+          }
+        }
+
+        PyPsiUtils.assertValid(typeComment);
+        typeComment.delete();
       }
     }
   }
@@ -100,21 +112,43 @@ public class PyConvertTypeCommentToVariableAnnotation extends PyBaseIntentionAct
 
   private static boolean isSuitableTypeComment(@NotNull PsiComment comment) {
     final String annotation = PyTypingTypeProvider.getTypeCommentValue(comment.getText());
-    return annotation != null && (findAssignmentTypeCommentTarget(comment) != null ||
-                                  findForLoopTypeCommentTarget(comment) != null ||
-                                  findWithStatementTypeCommentTarget(comment) != null);
+    return annotation != null && mapTargetsToAnnotations(comment) != null;
+  }
+
+  public static void runWithDocumentReleasedAndCommitted(@NotNull Project project, @NotNull Document document, @NotNull Runnable runnable) {
+    final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
+    manager.doPostponedOperationsAndUnblockDocument(document);
+    try {
+      runnable.run();
+    }
+    finally {
+      manager.commitDocument(document);
+    }
   }
 
   @Nullable
-  private static PyTargetExpression findAssignmentTypeCommentTarget(@NotNull PsiComment comment) {
-    final PsiElement parent = comment.getParent();
+  private static Map<PyTargetExpression, String> mapTargetsToAnnotations(@NotNull PsiComment typeComment) {
+    final PsiElement parent = typeComment.getParent();
     if (parent instanceof PyAssignmentStatement) {
       final PyAssignmentStatement assignment = (PyAssignmentStatement)parent;
       final PyExpression[] rawTargets = assignment.getRawTargets();
-      if (rawTargets.length == 1 && rawTargets[0] instanceof PyTargetExpression) {
-        final PyTargetExpression target = (PyTargetExpression)rawTargets[0];
-        if (target.getTypeComment() == comment) {
-          return target;
+      if (rawTargets.length == 1) {
+        return mapTargetsToAnnotations(rawTargets[0], typeComment);
+      }
+    }
+    else if (parent instanceof PyForPart) {
+      final PyForPart forPart = (PyForPart)parent;
+      final PyExpression target = forPart.getTarget();
+      if (target != null) {
+        return mapTargetsToAnnotations(target, typeComment);
+      }
+    }
+    else if (parent instanceof PyWithStatement) {
+      final PyWithItem[] withItems = ((PyWithStatement)parent).getWithItems();
+      if (withItems.length == 1) {
+        final PyExpression target = withItems[0].getTarget();
+        if (target != null) {
+          return mapTargetsToAnnotations(target, typeComment);
         }
       }
     }
@@ -122,44 +156,81 @@ public class PyConvertTypeCommentToVariableAnnotation extends PyBaseIntentionAct
   }
 
   @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;
+  private static Map<PyTargetExpression, String> mapTargetsToAnnotations(@NotNull PyExpression targetExpr, @NotNull PsiComment typeComment) {
+    final PyTargetExpression firstTarget = PsiTreeUtil.findChildOfType(targetExpr, PyTargetExpression.class, false);
+    if (firstTarget == null || firstTarget.getTypeComment() != typeComment) {
+      return null;
+    }
+
+    final String annotation = PyTypingTypeProvider.getTypeCommentValue(typeComment.getText());
+    if (annotation != null) {
+      if (targetExpr instanceof PyTargetExpression) {
+        return ImmutableMap.of((PyTargetExpression)targetExpr, annotation);
+      }
+
+      final PyElementGenerator generator = PyElementGenerator.getInstance(targetExpr.getProject());
+      final PyExpression parsed = generator.createExpressionFromText(LanguageLevel.PYTHON36, annotation);
+      if (parsed != null) {
+        return mapTargetsToAnnotations(targetExpr, parsed);
       }
     }
     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;
-        }
-      }
+  private static Map<PyTargetExpression, String> mapTargetsToAnnotations(@NotNull PyExpression targetExpr,
+                                                                         @NotNull PyExpression typeExpr) {
+    final PyExpression targetsNoParen = PyPsiUtils.flattenParens(targetExpr);
+    final PyExpression typesNoParen = PyPsiUtils.flattenParens(typeExpr);
+    if (targetsNoParen == null || typesNoParen == null) {
+      return null;
+    }
+    if (targetsNoParen instanceof PySequenceExpression && typesNoParen instanceof PySequenceExpression) {
+      final Ref<Map<PyTargetExpression, String>> result = new Ref<>(new LinkedHashMap<>());
+      mapTargetsToExpressions((PySequenceExpression)targetsNoParen, (PySequenceExpression)typesNoParen, result);
+      return result.get();
+    }
+    else if (targetsNoParen instanceof PyTargetExpression && !(typesNoParen instanceof PySequenceExpression)) {
+      return ImmutableMap.of((PyTargetExpression)targetsNoParen, typesNoParen.getText());
     }
     return null;
   }
 
+  private static void mapTargetsToExpressions(@NotNull PySequenceExpression targetSequence,
+                                              @NotNull PySequenceExpression valueSequence,
+                                              @NotNull Ref<Map<PyTargetExpression, String>> result) {
+    final PyExpression[] targets = targetSequence.getElements();
+    final PyExpression[] values = valueSequence.getElements();
 
-
-  public static void runWithDocumentReleasedAndCommitted(@NotNull Project project, @NotNull Document document, @NotNull Runnable runnable) {
-    final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
-    manager.doPostponedOperationsAndUnblockDocument(document);
-    try {
-      runnable.run();
+    if (targets.length != values.length) {
+      result.set(null);
+      return;
     }
-    finally {
-      manager.commitDocument(document);
+
+    for (int i = 0; i < targets.length; i++) {
+      final PyExpression target = PyPsiUtils.flattenParens(targets[i]);
+      final PyExpression value = PyPsiUtils.flattenParens(values[i]);
+
+      if (target == null || value == null) {
+        result.set(null);
+        return;
+      }
+
+      if (target instanceof PySequenceExpression && value instanceof PySequenceExpression) {
+        mapTargetsToExpressions((PySequenceExpression)target, (PySequenceExpression)value, result);
+        if (result.isNull()) {
+          return;
+        }
+      }
+      else if (target instanceof PyTargetExpression && !(value instanceof PySequenceExpression)) {
+        final Map<PyTargetExpression, String> map = result.get();
+        assert map != null;
+        map.put((PyTargetExpression)target, value.getText());
+      }
+      else {
+        result.set(null);
+        return;
+      }
     }
   }
 }
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithComplexUnpacking.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithComplexUnpacking.py
new file mode 100644 (file)
index 0000000..11c3f62
--- /dev/null
@@ -0,0 +1 @@
+[y, (x, (z))] = undefined()  # ty<caret>pe: Optional[Union[None, Any]], (Callable[..., int], Any)
\ No newline at end of file
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithComplexUnpacking_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithComplexUnpacking_after.py
new file mode 100644 (file)
index 0000000..f8bcc5c
--- /dev/null
@@ -0,0 +1,4 @@
+y: Optional[Union[None, Any]]
+x: Callable[..., int]
+z: Any
+[y, (x, (z))] = undefined()
\ No newline at end of file
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithUnpacking_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithUnpacking_after.py
new file mode 100644 (file)
index 0000000..2a16780
--- /dev/null
@@ -0,0 +1,3 @@
+x: int
+y: int
+x, y = (1, 2)
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithComplexUnpacking.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithComplexUnpacking.py
new file mode 100644 (file)
index 0000000..3391c71
--- /dev/null
@@ -0,0 +1,3 @@
+for ([y, (x, (z))]) in \
+        undefined():  # ty<caret>pe: (Tuple[int, ...], (Tuple[int], Tuple[Union[int, str]]))
+    pass
\ No newline at end of file
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithComplexUnpacking_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithComplexUnpacking_after.py
new file mode 100644 (file)
index 0000000..1042949
--- /dev/null
@@ -0,0 +1,6 @@
+y: Tuple[int, ...]
+x: Tuple[int]
+z: Tuple[Union[int, str]]
+for ([y, (x, (z))]) in \
+        undefined():
+    pass
\ No newline at end of file
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking_after.py
new file mode 100644 (file)
index 0000000..357a24a
--- /dev/null
@@ -0,0 +1,4 @@
+x: int
+y: str
+for x, y in undefined():
+    pass
\ No newline at end of file
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithComplexUnpacking.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithComplexUnpacking.py
new file mode 100644 (file)
index 0000000..eef06a5
--- /dev/null
@@ -0,0 +1,3 @@
+with undefined() \
+        as ((x, (z)), y):  # ty<caret>pe: (io.FileIO, Optional[int]), Any
+    pass
\ No newline at end of file
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithComplexUnpacking_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithComplexUnpacking_after.py
new file mode 100644 (file)
index 0000000..cfd29c2
--- /dev/null
@@ -0,0 +1,6 @@
+x: io.FileIO
+z: Optional[int]
+y: Any
+with undefined() \
+        as ((x, (z)), y):
+    pass
\ No newline at end of file
diff --git a/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking_after.py b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking_after.py
new file mode 100644 (file)
index 0000000..5fdd6f7
--- /dev/null
@@ -0,0 +1,4 @@
+x: int
+y: str
+with undefined() as (x, y):
+    pass
\ No newline at end of file
index 2bac1e73be14a8cf2b553e902089a3cefff15a53..a87770f5517e3329bd116b75a8e76f01f45b06a2 100644 (file)
@@ -44,7 +44,11 @@ public class PyConvertTypeCommentToVariableAnnotationIntentionTest extends PyInt
   }
 
   public void testAssignmentWithUnpacking() {
-    doNegativeTest();
+    doPositiveTest();
+  }
+
+  public void testAssignmentWithComplexUnpacking() {
+    doPositiveTest();
   }
 
   public void testMultilineAssignment() {
@@ -60,11 +64,19 @@ public class PyConvertTypeCommentToVariableAnnotationIntentionTest extends PyInt
   }
 
   public void testForLoopWithUnpacking() {
-    doNegativeTest();
+    doPositiveTest();
+  }
+
+  public void testForLoopWithComplexUnpacking() {
+    doPositiveTest();
   }
 
   public void testWithStatementWithUnpacking() {
-    doNegativeTest();
+    doPositiveTest();
+  }
+
+  public void testWithStatementWithComplexUnpacking() {
+    doPositiveTest();
   }
 
   public void testWithStatementWithMultipleWithItems() {