diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java index 9c687ee970c1..b1b7afadd524 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java @@ -116,37 +116,38 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { return candidates.anyMatch(PyAnnotateVariableTypeIntention::hasInlineAnnotation); } } - else { - final PyClass pyClass = target.getContainingClass(); - if (pyClass != null && scopeOwner instanceof PyFunction) { - final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); + else if (isInstanceAttribute(target, context)) { + // Set isDefinition=true to start searching right from the class level. + //noinspection ConstantConditions + final PyClassTypeImpl classType = new PyClassTypeImpl(target.getContainingClass(), true); + final List classAttrs = + classType.resolveMember(name, target, AccessDirection.READ, PyResolveContext.noImplicits().withTypeEvalContext(context), true); + if (classAttrs == null) { + return false; + } + return StreamEx.of(classAttrs) + .map(RatedResolveResult::getElement) + .select(PyTargetExpression.class) + .filter(x -> ScopeUtil.getScopeOwner(x) instanceof PyClass) + .anyMatch(PyAnnotateVariableTypeIntention::hasInlineAnnotation); + } + return false; + } - final boolean isInstanceAttribute; - if (context.maySwitchToAST(target)) { - //noinspection ConstantConditions - isInstanceAttribute = StreamEx.of(PyUtil.multiResolveTopPriority(target.getQualifier(), resolveContext)) - .select(PyParameter.class) - .filter(PyParameter::isSelf) - .anyMatch(p -> PsiTreeUtil.getParentOfType(p, PyFunction.class) == scopeOwner); - } - else { - isInstanceAttribute = PyUtil.isInstanceAttribute(target); - } - if (!isInstanceAttribute) { - return false; - } - // Set isDefinition=true to start searching right from the class level. - final PyClassTypeImpl classType = new PyClassTypeImpl(pyClass, true); - final List classAttrs = - classType.resolveMember(name, target, AccessDirection.READ, resolveContext, true); - if (classAttrs == null) { - return false; - } - return StreamEx.of(classAttrs) - .map(RatedResolveResult::getElement) - .select(PyTargetExpression.class) - .filter(x -> ScopeUtil.getScopeOwner(x) instanceof PyClass) - .anyMatch(PyAnnotateVariableTypeIntention::hasInlineAnnotation); + private static boolean isInstanceAttribute(@NotNull PyTargetExpression target, @NotNull TypeEvalContext context) { + final ScopeOwner scopeOwner = ScopeUtil.getScopeOwner(target); + if (target.isQualified() && target.getContainingClass() != null && scopeOwner instanceof PyFunction) { + + if (context.maySwitchToAST(target)) { + final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); + //noinspection ConstantConditions + return StreamEx.of(PyUtil.multiResolveTopPriority(target.getQualifier(), resolveContext)) + .select(PyParameter.class) + .filter(PyParameter::isSelf) + .anyMatch(p -> PsiTreeUtil.getParentOfType(p, PyFunction.class) == scopeOwner); + } + else { + return PyUtil.isInstanceAttribute(target); } } return false; @@ -177,7 +178,12 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { final TypeEvalContext context = TypeEvalContext.userInitiated(target.getProject(), target.getContainingFile()); final PyType inferredType = context.getType(target); final String annotationText = PythonDocumentationProvider.getTypeName(inferredType, context); - PyTypeHintGenerationUtil.insertVariableAnnotation(target, annotationText, true); + if (isInstanceAttribute(target, context)) { + PyTypeHintGenerationUtil.insertAttributeAnnotation(target, annotationText, true); + } + else { + PyTypeHintGenerationUtil.insertVariableAnnotation(target, annotationText, true); + } } private static void insertVariableTypeComment(@NotNull PyTargetExpression target) { diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java index f79183fb9fe9..fba0689a346d 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java @@ -19,6 +19,7 @@ import com.intellij.psi.util.PsiTreeUtil; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyPsiUtils; import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; import java.util.Collections; import java.util.List; @@ -31,6 +32,39 @@ import static com.jetbrains.python.psi.PyUtil.as; public class PyTypeHintGenerationUtil { private PyTypeHintGenerationUtil() {} + public static void insertAttributeAnnotation(@NotNull PyTargetExpression target, @NotNull String annotation, boolean startTemplate) { + final LanguageLevel langLevel = LanguageLevel.forElement(target); + if (langLevel.isOlderThan(LanguageLevel.PYTHON36)) { + throw new IllegalArgumentException("Target '" + target.getText() + "' doesn't belong to Python 3.6+ project: " + langLevel); + } + + 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 String declarationText = target.getName() + ": " + annotation; + final PyTypeDeclarationStatement declaration = generator.createFromText(langLevel, PyTypeDeclarationStatement.class, declarationText); + final PsiElement anchorBefore = findPrecedingAnchorForAttributeDeclaration(pyClass); + PyTypeDeclarationStatement inserted = (PyTypeDeclarationStatement)pyClass.getStatementList().addAfter(declaration, anchorBefore); + inserted = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(inserted); + + if (startTemplate && inserted != null) { + openEditorAndAddTemplateForAnnotation(inserted); + } + } + + @Nullable + private static PsiElement findPrecedingAnchorForAttributeDeclaration(@NotNull PyClass pyClass) { + final PyStatement firstStatement = pyClass.getStatementList().getStatements()[0]; + final PyStringLiteralExpression classDocstring = pyClass.getDocStringExpression(); + if (firstStatement instanceof PyExpressionStatement && classDocstring == ((PyExpressionStatement)firstStatement).getExpression()) { + return firstStatement; + } + return null; + } + public static void insertVariableAnnotation(@NotNull PyTargetExpression target, @NotNull String annotation, boolean startTemplate) { final LanguageLevel langLevel = LanguageLevel.forElement(target); if (langLevel.isOlderThan(LanguageLevel.PYTHON36)) { @@ -50,7 +84,7 @@ public class PyTypeHintGenerationUtil { } else { final PyElementGenerator generator = PyElementGenerator.getInstance(project); - final String declarationText = target.getText() + ": " + annotation; + final String declarationText = target.getName() + ": " + annotation; final PyTypeDeclarationStatement declaration = generator.createFromText(langLevel, PyTypeDeclarationStatement.class, declarationText); final PyStatement statement = PsiTreeUtil.getParentOfType(target, PyStatement.class); assert statement != null; @@ -59,21 +93,28 @@ public class PyTypeHintGenerationUtil { } if (startTemplate && createdAnnotationOwner != null) { - assert createdAnnotationOwner.getAnnotationValue() != null; + openEditorAndAddTemplateForAnnotation(createdAnnotationOwner); + } + } - final int initialCaretOffset = createdAnnotationOwner.getTextRange().getStartOffset(); - final VirtualFile updatedVirtualFile = createdAnnotationOwner.getContainingFile().getVirtualFile(); - final OpenFileDescriptor descriptor = new OpenFileDescriptor(project, updatedVirtualFile, initialCaretOffset); - final Editor editor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); + private static void openEditorAndAddTemplateForAnnotation(@NotNull PyAnnotationOwner annotated) { + assert annotated.isValid(); + assert annotated.getAnnotationValue() != null; - if (editor != null) { - editor.getCaretModel().moveToOffset(initialCaretOffset); - final TemplateBuilder templateBuilder = TemplateBuilderFactory.getInstance().createTemplateBuilder(createdAnnotationOwner); - final String replacementText = ApplicationManager.getApplication().isUnitTestMode() ? "[" + annotation + "]" : annotation; - //noinspection ConstantConditions - templateBuilder.replaceElement(createdAnnotationOwner.getAnnotation().getValue(), replacementText); - templateBuilder.run(editor, true); - } + final Project project = annotated.getProject(); + final int initialCaretOffset = annotated.getTextRange().getStartOffset(); + final VirtualFile updatedVirtualFile = annotated.getContainingFile().getVirtualFile(); + final OpenFileDescriptor descriptor = new OpenFileDescriptor(project, updatedVirtualFile, initialCaretOffset); + final Editor editor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); + + if (editor != null) { + editor.getCaretModel().moveToOffset(initialCaretOffset); + final TemplateBuilder templateBuilder = TemplateBuilderFactory.getInstance().createTemplateBuilder(annotated); + final String annotation = annotated.getAnnotationValue(); + final String replacementText = ApplicationManager.getApplication().isUnitTestMode() ? "[" + annotation + "]" : annotation; + //noinspection ConstantConditions + templateBuilder.replaceElement(annotated.getAnnotation().getValue(), replacementText); + templateBuilder.run(editor, true); } } diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttribute.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttribute.py new file mode 100644 index 000000000000..6c116981cbe6 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttribute.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/annotationInstanceAttributeDocstring.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeDocstring.py new file mode 100644 index 000000000000..8a413fe9f931 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeDocstring.py @@ -0,0 +1,6 @@ +class MyClass: + """Docstring.""" + + def __init__(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeDocstring_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeDocstring_after.py new file mode 100644 index 000000000000..401f04c4f7f2 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeDocstring_after.py @@ -0,0 +1,7 @@ +class MyClass: + """Docstring.""" + attr: [int] + + def __init__(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttribute_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttribute_after.py new file mode 100644 index 000000000000..40be05ad0683 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttribute_after.py @@ -0,0 +1,6 @@ +class MyClass: + attr: [int] + + def __init__(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedInstanceAttributeWithAnnotation.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedInstanceAttributeWithAnnotation.py new file mode 100644 index 000000000000..0adf21a2f9fd --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedInstanceAttributeWithAnnotation.py @@ -0,0 +1,6 @@ +class MyClass: + attr: int + + def __init___(self, x): + self.attr = x + self.attr diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedInstanceAttributeWithTypeComment.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedInstanceAttributeWithTypeComment.py new file mode 100644 index 000000000000..be63f1f2e369 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedInstanceAttributeWithTypeComment.py @@ -0,0 +1,6 @@ +class MyClass: + attr = ... # type: int + + def __init___(self, x): + self.attr = x + self.attr diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttribute.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttribute.py new file mode 100644 index 000000000000..6c116981cbe6 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttribute.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/typeCommentInstanceAttributeDocstring.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring.py new file mode 100644 index 000000000000..8a413fe9f931 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring.py @@ -0,0 +1,6 @@ +class MyClass: + """Docstring.""" + + def __init__(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring_after.py new file mode 100644 index 000000000000..354a85cb441a --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeDocstring_after.py @@ -0,0 +1,7 @@ +class MyClass: + """Docstring.""" + attr = None # type: [int] + + def __init__(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttribute_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttribute_after.py new file mode 100644 index 000000000000..9dc70694b6d2 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttribute_after.py @@ -0,0 +1,6 @@ +class MyClass: + attr = None # 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 b0b2fbf2b328..86db74c48984 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java @@ -50,6 +50,14 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { doNegativeTest(); } + public void testNotSuggestedInstanceAttributeWithAnnotation() { + doNegativeTest(); + } + + public void testNotSuggestedInstanceAttributeWithTypeComment() { + doNegativeTest(); + } + public void testAnnotationLocalSimpleAssignmentTarget() { doTestAnnotation(); } @@ -74,6 +82,22 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { doTestAnnotation(); } + public void testAnnotationInstanceAttribute() { + doTestAnnotation(); + } + + public void testAnnotationInstanceAttributeDocstring() { + doTestAnnotation(); + } + + public void testTypeCommentInstanceAttribute() { + doTestAnnotation(); + } + + public void testTypeCommentInstanceAttributeDocstring() { + doTestAnnotation(); + } + public void testTypeCommentLocalSimpleAssignmentTarget() { doTestTypeComment(); }