From 9bbbbc6eab16d8caf8112fae3dceadb495cae617 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Thu, 8 Mar 2018 19:43:37 +0300 Subject: [PATCH] PY-21191 PY-28879 Support nested tuples in type hints inside comments by reusing the logic from "Convert to variable annotation" intention This change also partially remedies the problem described in PY-28879 when the whole file was parsed whenever we found typing.Tuple in a variable type comment, because we no longer transform type hints like "int, str" into "Tuple[int, str]" as an intermediate step and process them right away instead. We still need to handle general type comments with unpacking in stub-safe manner, though. --- ...eCommentToVariableAnnotationIntention.java | 77 ++---------- .../typing/PyTypingTypeProvider.java | 111 +++++++++++++++--- .../stubs/VariableTypeCommentWithTupleType.py | 3 + .../com/jetbrains/python/PyStubsTest.java | 10 +- .../com/jetbrains/python/PyTypingTest.java | 38 ++++++ 5 files changed, 155 insertions(+), 84 deletions(-) create mode 100644 python/testData/stubs/VariableTypeCommentWithTupleType.py diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java index bdc08126d67a..a3a7abc18f4d 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertTypeCommentToVariableAnnotationIntention.java @@ -18,7 +18,6 @@ package com.jetbrains.python.codeInsight.intentions; import com.google.common.collect.ImmutableMap; 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.PsiElement; import com.intellij.psi.PsiFile; @@ -31,11 +30,13 @@ import com.jetbrains.python.codeInsight.intentions.PyTypeHintGenerationUtil.Anno import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyPsiUtils; +import one.util.streamex.EntryStream; import org.jetbrains.annotations.Nls; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import java.util.ArrayList; +import java.util.Collections; import java.util.Map; public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseIntentionAction { @@ -44,7 +45,7 @@ public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseInt final PsiComment typeComment = findCommentUnderCaret(editor, file); if (typeComment != null) { final Map map = mapTargetsToAnnotations(typeComment); - if (map != null) { + if (!map.isEmpty()) { if (typeComment.getParent() instanceof PyAssignmentStatement && map.size() == 1) { final PyTargetExpression target = ContainerUtil.getFirstItem(map.keySet()); assert target != null; @@ -92,10 +93,10 @@ public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseInt private static boolean isSuitableTypeComment(@NotNull PsiComment comment) { final String annotation = PyTypingTypeProvider.getTypeCommentValue(comment.getText()); - return annotation != null && mapTargetsToAnnotations(comment) != null; + return annotation != null && !mapTargetsToAnnotations(comment).isEmpty(); } - @Nullable + @NotNull private static Map mapTargetsToAnnotations(@NotNull PsiComment typeComment) { final PsiElement parent = typeComment.getParent(); if (parent instanceof PyAssignmentStatement) { @@ -121,14 +122,14 @@ public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseInt } } } - return null; + return Collections.emptyMap(); } - @Nullable + @NotNull 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; + return Collections.emptyMap(); } final String annotation = PyTypingTypeProvider.getTypeCommentValue(typeComment.getText()); @@ -140,66 +141,10 @@ public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseInt 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; + final Map targetToExpr = PyTypingTypeProvider.mapTargetsToAnnotations(targetExpr, parsed); + return EntryStream.of(targetToExpr).mapValues(PyExpression::getText).toCustomMap(LinkedHashMap::new); } } + return Collections.emptyMap(); } } diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java index 36d02f9a9b09..f2e5d2f6c0f1 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java @@ -641,17 +641,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { } final String comment = target.getTypeCommentAnnotation(); if (comment != null) { - final Ref fromTypeComment = getVariableTypeCommentType(comment, target, new Context(context)); - if (fromTypeComment != null) { - final PyType type = Ref.deref(fromTypeComment); - if (type instanceof PyTupleType) { - final PyTupleExpression tupleExpr = PsiTreeUtil.getParentOfType(target, PyTupleExpression.class); - if (tupleExpr != null) { - return Ref.create(PyTypeChecker.getTargetTypeFromTupleAssignment(target, tupleExpr, (PyTupleType)type)); - } - } - return fromTypeComment; - } + return getVariableTypeCommentType(comment, target, new Context(context)); } return null; } @@ -1004,21 +994,108 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { } @Nullable - private static Ref getVariableTypeCommentType(@NotNull String contents, @NotNull PsiElement anchor, @NotNull Context context) { + private static Ref getVariableTypeCommentType(@NotNull String contents, + @NotNull PyTargetExpression target, + @NotNull Context context) { // TODO pass the real anchor as the context element for the fragment to resolve local classes/type aliases - final PyExpression expr = PyPsiUtils.flattenParens(PyUtil.createExpressionFromFragment(contents, anchor.getContainingFile())); + final PyExpression expr = PyPsiUtils.flattenParens(PyUtil.createExpressionFromFragment(contents, target.getContainingFile())); if (expr != null) { // Such syntax is specific to "# type:" comments, unpacking in type hints is not allowed anywhere else if (expr instanceof PyTupleExpression) { - final PyTupleExpression tupleExpr = (PyTupleExpression)expr; - final List elementTypes = ContainerUtil.map(tupleExpr.getElements(), elementExpr -> Ref.deref(getType(elementExpr, context))); - return Ref.create(PyTupleType.create(anchor, elementTypes)); + // XXX: Switches stub to AST + final PyExpression topmostTarget = findTopmostTarget(target); + if (topmostTarget != null) { + final Map targetToExpr = mapTargetsToAnnotations(topmostTarget, expr); + final PyExpression typeExpr = targetToExpr.get(target); + if (typeExpr != null) { + return getType(typeExpr, context); + } + } + } + else { + return getType(expr, context); } - return getType(expr, context); } return null; } + @Nullable + private static PyExpression findTopmostTarget(@NotNull PyTargetExpression target) { + final PyElement validTargetParent = PsiTreeUtil.getParentOfType(target, PyForPart.class, PyWithItem.class, PyAssignmentStatement.class); + if (validTargetParent == null) { + return null; + } + final PyExpression topmostTarget = as(PsiTreeUtil.findPrevParent(validTargetParent, target), PyExpression.class); + if (validTargetParent instanceof PyForPart && topmostTarget != ((PyForPart)validTargetParent).getTarget()) { + return null; + } + if (validTargetParent instanceof PyWithItem && topmostTarget != ((PyWithItem)validTargetParent).getTarget()) { + return null; + } + if (validTargetParent instanceof PyAssignmentStatement && + ArrayUtil.indexOf(((PyAssignmentStatement)validTargetParent).getRawTargets(), topmostTarget) < 0) { + return null; + } + return topmostTarget; + } + + @NotNull + public 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 Collections.emptyMap(); + } + if (targetsNoParen instanceof PySequenceExpression && typesNoParen instanceof PySequenceExpression) { + final Ref> result = new Ref<>(new LinkedHashMap<>()); + mapTargetsToExpressions((PySequenceExpression)targetsNoParen, (PySequenceExpression)typesNoParen, result); + return result.isNull() ? Collections.emptyMap() : Collections.unmodifiableMap(result.get()); + } + else if (targetsNoParen instanceof PyTargetExpression && !(typesNoParen instanceof PySequenceExpression)) { + return ImmutableMap.of((PyTargetExpression)targetsNoParen, typesNoParen); + } + return Collections.emptyMap(); + } + + 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); + } + else { + result.set(null); + return; + } + } + } + @Nullable private static PyType getCallableType(@NotNull PsiElement resolved, @NotNull Context context) { if (resolved instanceof PySubscriptionExpression) { diff --git a/python/testData/stubs/VariableTypeCommentWithTupleType.py b/python/testData/stubs/VariableTypeCommentWithTupleType.py new file mode 100644 index 000000000000..5683d946aaca --- /dev/null +++ b/python/testData/stubs/VariableTypeCommentWithTupleType.py @@ -0,0 +1,3 @@ +from typing import Tuple + +var = None # type: Tuple[int, str] diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java index 9fe26f47e5b3..f2d61e57856c 100644 --- a/python/testSrc/com/jetbrains/python/PyStubsTest.java +++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java @@ -2,7 +2,6 @@ package com.jetbrains.python; import com.intellij.lang.FileASTNode; -import com.intellij.openapi.application.Result; import com.intellij.openapi.command.WriteCommandAction; import com.intellij.openapi.editor.Document; import com.intellij.openapi.project.DumbServiceImpl; @@ -983,4 +982,13 @@ public class PyStubsTest extends PyTestCase { } ); } + + // PY-28879 + public void testVariableTypeCommentWithTupleType() { + final PyFile file = getTestFile(); + final PyTargetExpression target = file.findTopLevelAttribute("var"); + final TypeEvalContext context = TypeEvalContext.codeInsightFallback(target.getProject()); + context.getType(target); + assertNotParsed(file); + } } diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index 347c4dc1553b..ee816fdcad4c 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java @@ -350,6 +350,44 @@ public class PyTypingTest extends PyTestCase { "expr, x = undefined() # type: (int, str) "); } + // PY-21191 + public void testTypeCommentWithNestedTuplesInAssignment() { + doTest("int", + "_, (_, expr) = undefined() # type: str, (str, int)"); + } + + // PY-21191 + public void testTypeCommentStructuralMismatch1() { + doTest("Any", + "expr = undefined() # type: str, int"); + } + + // PY-21191 + public void testTypeCommentStructuralMismatch2() { + doTest("Any", + "_, (_, expr) = undefined() # type: str, (str, str, int)"); + } + + // PY-21191 + public void testTypeCommentStructuralMismatch3() { + doTest("Any", + "_, (_, expr) = undefined() # type: (str, str), int"); + } + + // PY-21191 + public void testTypeCommentWithNestedTuplesInWithStatement() { + doTest("int", + "with undefined() as (_, (_, expr)): # type: str, (str, int)\n" + + " pass"); + } + + // PY-21191 + public void testTypeCommentWithNestedTuplesInForStatement() { + doTest("int", + "for (_, (_, expr)) in undefined(): # type: str, (str, int)\n" + + " pass"); + } + // PY-16585 public void testCommentAfterComprehensionInAssignment() { doTest("int",