From 3c74966e9535724985f4a1f29239b390967fe18d Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Mon, 22 Jan 2018 21:11:03 +0300 Subject: [PATCH] PY-16412 Support type hints for instance attributes using type comments --- .../PyAnnotateVariableTypeIntention.java | 13 ++-- .../intentions/PyTypeHintGenerationUtil.java | 72 ++++++++++++++----- ...CommentInstanceAttributeDocstring_after.py | 2 +- .../typeCommentInstanceAttributePy3.py | 4 ++ .../typeCommentInstanceAttributePy3_after.py | 6 ++ .../PyAnnotateVariableTypeIntentionTest.java | 8 ++- 6 files changed, 79 insertions(+), 26 deletions(-) create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributePy3.py create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributePy3_after.py diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java index b1b7afadd524..e44fafee6937 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java @@ -187,13 +187,18 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { } private static void insertVariableTypeComment(@NotNull PyTargetExpression target) { - final Pair> annotationAndRanges = generateNestedTypeHint(target); - PyTypeHintGenerationUtil.insertVariableTypeComment(target, annotationAndRanges.getFirst(), true, annotationAndRanges.getSecond()); + final TypeEvalContext context = TypeEvalContext.userInitiated(target.getProject(), target.getContainingFile()); + final Pair> annotationAndRanges = generateNestedTypeHint(target, context); + if (isInstanceAttribute(target, context)) { + PyTypeHintGenerationUtil.insertAttributeTypeComment(target, annotationAndRanges.getFirst(), true, annotationAndRanges.getSecond()); + } + else { + PyTypeHintGenerationUtil.insertVariableTypeComment(target, annotationAndRanges.getFirst(), true, annotationAndRanges.getSecond()); + } } @NotNull - private static Pair> generateNestedTypeHint(@NotNull PyTargetExpression target) { - final TypeEvalContext context = TypeEvalContext.userInitiated(target.getProject(), target.getContainingFile()); + private static Pair> generateNestedTypeHint(@NotNull PyTargetExpression target, TypeEvalContext context) { final PyElement validTargetParent = PsiTreeUtil.getParentOfType(target, PyForPart.class, PyWithItem.class, PyAssignmentStatement.class); assert validTargetParent != null; final PsiElement topmostTarget = PsiTreeUtil.findPrevParent(validTargetParent, target); diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java index fba0689a346d..412fc5a4d202 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java @@ -30,8 +30,37 @@ import static com.jetbrains.python.psi.PyUtil.as; * @author Mikhail Golubev */ public class PyTypeHintGenerationUtil { + + public static final String TYPE_COMMENT_PREFIX = "# type: "; + private PyTypeHintGenerationUtil() {} + public static void insertAttributeTypeComment(@NotNull PyTargetExpression target, + @NotNull String annotation, + boolean startTemplate, + @NotNull List typeRanges) { + + final PyClass pyClass = target.getContainingClass(); + if (pyClass == null) { + throw new IllegalArgumentException("Target '" + target.getText() + "' in not contained in a class definition"); + } + + final PyElementGenerator generator = PyElementGenerator.getInstance(target.getProject()); + final LanguageLevel langLevel = LanguageLevel.forElement(target); + final String assignedValue = langLevel.isAtLeast(LanguageLevel.PYTHON30) ? "..." : "None"; + final String declarationText = target.getName() + " = " + assignedValue + " " + TYPE_COMMENT_PREFIX + annotation; + final PyAssignmentStatement declaration = generator.createFromText(langLevel, PyAssignmentStatement.class, declarationText); + final PsiElement anchorBefore = findPrecedingAnchorForAttributeDeclaration(pyClass); + PyAssignmentStatement inserted = (PyAssignmentStatement)pyClass.getStatementList().addAfter(declaration, anchorBefore); + inserted = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(inserted); + + final PsiComment insertedComment = as(inserted.getLastChild(), PsiComment.class); + + if (startTemplate && insertedComment != null) { + openEditorAndAddTemplateForTypeComment(insertedComment, annotation, typeRanges); + } + } + public static void insertAttributeAnnotation(@NotNull PyTargetExpression target, @NotNull String annotation, boolean startTemplate) { final LanguageLevel langLevel = LanguageLevel.forElement(target); if (langLevel.isOlderThan(LanguageLevel.PYTHON36)) { @@ -127,15 +156,14 @@ public class PyTypeHintGenerationUtil { public static void insertVariableTypeComment(@NotNull PyTargetExpression target, @NotNull String annotation, boolean startTemplate) { - insertVariableTypeComment(target, annotation, startTemplate, Collections.singletonList(TextRange.from(0, annotation.length()))); + insertVariableTypeComment(target, annotation, startTemplate, Collections.singletonList(TextRange.allOf(annotation))); } public static void insertVariableTypeComment(@NotNull PyTargetExpression target, @NotNull String annotation, boolean startTemplate, @NotNull List typeRanges) { - final String typeCommentPrefix = "# type: "; - final String typeCommentText = " " + typeCommentPrefix + annotation; + final String typeCommentText = " " + TYPE_COMMENT_PREFIX + annotation; final PyStatement statement = PsiTreeUtil.getParentOfType(target, PyStatement.class); final PsiElement insertionAnchor; @@ -170,24 +198,30 @@ public class PyTypeHintGenerationUtil { final PsiComment insertedComment = target.getTypeComment(); if (startTemplate && insertedComment != null) { - final int initialCaretOffset = insertedComment.getTextRange().getStartOffset(); - final VirtualFile updatedVirtualFile = insertedComment.getContainingFile().getVirtualFile(); - final Project project = target.getProject(); - final OpenFileDescriptor descriptor = new OpenFileDescriptor(project, updatedVirtualFile, initialCaretOffset); - final Editor editor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); + openEditorAndAddTemplateForTypeComment(insertedComment, annotation, typeRanges); + } + } - if (editor != null) { - final boolean testMode = ApplicationManager.getApplication().isUnitTestMode(); - editor.getCaretModel().moveToOffset(initialCaretOffset); - final TemplateBuilder templateBuilder = TemplateBuilderFactory.getInstance().createTemplateBuilder(insertedComment); - //noinspection ConstantConditions - for (TextRange range : typeRanges) { - final String individualType = range.substring(annotation); - final String replacementText = testMode ? "[" + individualType + "]" : individualType; - templateBuilder.replaceRange(range.shiftRight(typeCommentPrefix.length()), replacementText); - } - templateBuilder.run(editor, true); + private static void openEditorAndAddTemplateForTypeComment(@NotNull PsiComment insertedComment, + @NotNull String annotation, + @NotNull List typeRanges) { + final int initialCaretOffset = insertedComment.getTextRange().getStartOffset(); + final VirtualFile updatedVirtualFile = insertedComment.getContainingFile().getVirtualFile(); + final Project project = insertedComment.getProject(); + final OpenFileDescriptor descriptor = new OpenFileDescriptor(project, updatedVirtualFile, initialCaretOffset); + final Editor editor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); + + if (editor != null) { + final boolean testMode = ApplicationManager.getApplication().isUnitTestMode(); + editor.getCaretModel().moveToOffset(initialCaretOffset); + final TemplateBuilder templateBuilder = TemplateBuilderFactory.getInstance().createTemplateBuilder(insertedComment); + //noinspection ConstantConditions + for (TextRange range : typeRanges) { + final String individualType = range.substring(annotation); + final String replacementText = testMode ? "[" + individualType + "]" : individualType; + templateBuilder.replaceRange(range.shiftRight(TYPE_COMMENT_PREFIX.length()), replacementText); } + templateBuilder.run(editor, true); } } } diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring_after.py index 354a85cb441a..cf0c4e297536 100644 --- a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring_after.py +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring_after.py @@ -4,4 +4,4 @@ class MyClass: def __init__(self): self.attr = 42 - self.attr \ No newline at end of file + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributePy3.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributePy3.py new file mode 100644 index 000000000000..6c116981cbe6 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributePy3.py @@ -0,0 +1,4 @@ +class MyClass: + def __init__(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributePy3_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributePy3_after.py new file mode 100644 index 000000000000..83479f63db36 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributePy3_after.py @@ -0,0 +1,6 @@ +class MyClass: + attr = ... # type: [int] + + def __init__(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java index 86db74c48984..2c9e5cb237b3 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java @@ -91,11 +91,15 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { } public void testTypeCommentInstanceAttribute() { - doTestAnnotation(); + doTestTypeComment(); + } + + public void testTypeCommentInstanceAttributePy3() { + doTest(LanguageLevel.PYTHON30); } public void testTypeCommentInstanceAttributeDocstring() { - doTestAnnotation(); + doTestTypeComment(); } public void testTypeCommentLocalSimpleAssignmentTarget() {