diff --git a/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/after.py.template b/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/after.py.template index 89023c6848ad..680d63c08393 100644 --- a/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/after.py.template +++ b/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/after.py.template @@ -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 diff --git a/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/before.py.template b/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/before.py.template index e28b15f6074b..4af9ffca879b 100644 --- a/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/before.py.template +++ b/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/before.py.template @@ -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 + ... diff --git a/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/description.html b/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/description.html index 2d034cc57781..6d235944ba5e 100644 --- a/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/description.html +++ b/python/resources/intentionDescriptions/PyConvertTypeCommentToVariableAnnotation/description.html @@ -1,7 +1,7 @@ - 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. \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java index e8e2d7fc8cde..9bb6485bc425 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotation.java @@ -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 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> entries = new ArrayList<>(map.entrySet()); + Collections.reverse(entries); + + PsiElement anchor = statement; + for (Map.Entry 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 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 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 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> 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> 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 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 index 000000000000..11c3f623cd34 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithComplexUnpacking.py @@ -0,0 +1 @@ +[y, (x, (z))] = undefined() # type: 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 index 000000000000..f8bcc5ccd725 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithComplexUnpacking_after.py @@ -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 index 000000000000..2a1678083c31 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/assignmentWithUnpacking_after.py @@ -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 index 000000000000..3391c7142747 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithComplexUnpacking.py @@ -0,0 +1,3 @@ +for ([y, (x, (z))]) in \ + undefined(): # type: (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 index 000000000000..104294980762 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithComplexUnpacking_after.py @@ -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 index 000000000000..357a24a6415e --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/forLoopWithUnpacking_after.py @@ -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 index 000000000000..eef06a5c6529 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithComplexUnpacking.py @@ -0,0 +1,3 @@ +with undefined() \ + as ((x, (z)), y): # type: (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 index 000000000000..cfd29c242b82 --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithComplexUnpacking_after.py @@ -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 index 000000000000..5fdd6f763a2a --- /dev/null +++ b/python/testData/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest/withStatementWithUnpacking_after.py @@ -0,0 +1,4 @@ +x: int +y: str +with undefined() as (x, y): + 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 2bac1e73be14..a87770f5517e 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyConvertTypeCommentToVariableAnnotationIntentionTest.java @@ -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() {