PY-20801 Handle cases with nested target expressions and unpacking

2026-09-27 10:03:11 +07:00 · 2016-10-25 00:03:49 +03:00
parent 5e772eb11b
commit c1e7e66edf
14 changed files with 221 additions and 85 deletions
@@ -1 +1,13 @@
-x: Optional[int] = undefined()
+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:
+    ...
@@ -1 +1,8 @@
-x = undefined()  # type: Optional[int]
+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
+    ...
@@ -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>
@@ -15,15 +15,20 @@
 */
 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,58 +112,9 @@ 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;
  }

-  @Nullable
-  private static PyTargetExpression findAssignmentTypeCommentTarget(@NotNull PsiComment comment) {
-    final PsiElement parent = comment.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;
-        }
-      }
-    }
-    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);
@@ -162,4 +125,112 @@ public class PyConvertTypeCommentToVariableAnnotation extends PyBaseIntentionAct
      manager.commitDocument(document);
    }
  }
+
+  @Nullable
+  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) {
+        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);
+        }
+      }
+    }
+    return null;
+  }
+
+  @Nullable
+  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 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();
+
+    if (targets.length != values.length) {
+      result.set(null);
+      return;
+    }
+
+    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;
+      }
+    }
+  }
 }
@@ -0,0 +1 @@
+[y, (x, (z))] = undefined()  # ty<caret>pe: Optional[Union[None, Any]], (Callable[..., int], Any)
@@ -0,0 +1,4 @@
+y: Optional[Union[None, Any]]
+x: Callable[..., int]
+z: Any
+[y, (x, (z))] = undefined()
@@ -0,0 +1,3 @@
+x: int
+y: int
+x, y = (1, 2)
@@ -0,0 +1,3 @@
+for ([y, (x, (z))]) in \
+        undefined():  # ty<caret>pe: (Tuple[int, ...], (Tuple[int], Tuple[Union[int, str]]))
+    pass
@@ -0,0 +1,6 @@
+y: Tuple[int, ...]
+x: Tuple[int]
+z: Tuple[Union[int, str]]
+for ([y, (x, (z))]) in \
+        undefined():
+    pass
@@ -0,0 +1,4 @@
+x: int
+y: str
+for x, y in undefined():
+    pass
@@ -0,0 +1,3 @@
+with undefined() \
+        as ((x, (z)), y):  # ty<caret>pe: (io.FileIO, Optional[int]), Any
+    pass
@@ -0,0 +1,6 @@
+x: io.FileIO
+z: Optional[int]
+y: Any
+with undefined() \
+        as ((x, (z)), y):
+    pass
@@ -0,0 +1,4 @@
+x: int
+y: str
+with undefined() as (x, y):
+    pass
@@ -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() {
				`@@ -0,0 +1 @@`
				`[y, (x, (z))] = undefined() # ty<caret>pe: Optional[Union[None, Any]], (Callable[..., int], Any)`