From bf5a0dd8b19ed80cd18296578a179264ea944385 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Tue, 23 Jan 2018 12:07:11 +0300 Subject: [PATCH] PY-16412 Add annotations and type comments to existing class-level definitions for instance attributes --- .../PyAnnotateVariableTypeIntention.java | 53 ++++++++++++++----- ...onInstanceAttributeClassLevelAssignment.py | 6 +++ ...AttributeClassLevelAssignmentInAncestor.py | 8 +++ ...uteClassLevelAssignmentInAncestor_after.py | 8 +++ ...anceAttributeClassLevelAssignment_after.py | 6 +++ ...ntInstanceAttributeClassLevelAssignment.py | 6 +++ ...AttributeClassLevelAssignmentInAncestor.py | 8 +++ ...uteClassLevelAssignmentInAncestor_after.py | 8 +++ ...anceAttributeClassLevelAssignment_after.py | 6 +++ .../PyAnnotateVariableTypeIntentionTest.java | 16 ++++++ 10 files changed, 111 insertions(+), 14 deletions(-) create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignment.py create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignmentInAncestor.py create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignmentInAncestor_after.py create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignment_after.py create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignment.py create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignmentInAncestor.py create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignmentInAncestor_after.py create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignment_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 e44fafee6937..507e1e3a1e67 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java @@ -10,6 +10,7 @@ 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.jetbrains.python.PyBundle; import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache; import com.jetbrains.python.codeInsight.controlflow.ScopeOwner; @@ -119,21 +120,30 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { 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); + final List classLevelDefinitions = findClassLevelDefinitions(target, context); + return ContainerUtil.exists(classLevelDefinitions, PyAnnotateVariableTypeIntention::hasInlineAnnotation); } return false; } + @NotNull + private static List findClassLevelDefinitions(@NotNull PyTargetExpression target, @NotNull TypeEvalContext context) { + assert target.getContainingClass() != null; + assert target.getName() != null; + final PyClassTypeImpl classType = new PyClassTypeImpl(target.getContainingClass(), true); + final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); + final List classAttrs = + classType.resolveMember(target.getName(), target, AccessDirection.READ, resolveContext, true); + if (classAttrs == null) { + return Collections.emptyList(); + } + return StreamEx.of(classAttrs) + .map(RatedResolveResult::getElement) + .select(PyTargetExpression.class) + .filter(x -> ScopeUtil.getScopeOwner(x) instanceof PyClass) + .toList(); + } + 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) { @@ -179,7 +189,13 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { final PyType inferredType = context.getType(target); final String annotationText = PythonDocumentationProvider.getTypeName(inferredType, context); if (isInstanceAttribute(target, context)) { - PyTypeHintGenerationUtil.insertAttributeAnnotation(target, annotationText, true); + final List classLevelAttrs = findClassLevelDefinitions(target, context); + if (classLevelAttrs.isEmpty()) { + PyTypeHintGenerationUtil.insertAttributeAnnotation(target, annotationText, true); + } + else { + PyTypeHintGenerationUtil.insertVariableAnnotation(classLevelAttrs.get(0), annotationText, true); + } } else { PyTypeHintGenerationUtil.insertVariableAnnotation(target, annotationText, true); @@ -189,11 +205,20 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { private static void insertVariableTypeComment(@NotNull PyTargetExpression target) { final TypeEvalContext context = TypeEvalContext.userInitiated(target.getProject(), target.getContainingFile()); final Pair> annotationAndRanges = generateNestedTypeHint(target, context); + final String annotationText = annotationAndRanges.getFirst(); + final List typeRanges = annotationAndRanges.getSecond(); if (isInstanceAttribute(target, context)) { - PyTypeHintGenerationUtil.insertAttributeTypeComment(target, annotationAndRanges.getFirst(), true, annotationAndRanges.getSecond()); + final List classLevelAttrs = findClassLevelDefinitions(target, context); + if (classLevelAttrs.isEmpty()) { + PyTypeHintGenerationUtil.insertAttributeTypeComment(target, annotationText, true, typeRanges); + } + else { + // Use existing class level definition (say, assignment of the default value) for annotation + PyTypeHintGenerationUtil.insertVariableTypeComment(classLevelAttrs.get(0), annotationText, true, typeRanges); + } } else { - PyTypeHintGenerationUtil.insertVariableTypeComment(target, annotationAndRanges.getFirst(), true, annotationAndRanges.getSecond()); + PyTypeHintGenerationUtil.insertVariableTypeComment(target, annotationText, true, typeRanges); } } diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignment.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignment.py new file mode 100644 index 000000000000..03f3eb1aad9c --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignment.py @@ -0,0 +1,6 @@ +class MyClass: + attr = 0 + + def __init___(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignmentInAncestor.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignmentInAncestor.py new file mode 100644 index 000000000000..46a47f7f0d59 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignmentInAncestor.py @@ -0,0 +1,8 @@ +class Base(object): + attr = 0 + + +class MyClass(Base): + def __init___(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignmentInAncestor_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignmentInAncestor_after.py new file mode 100644 index 000000000000..f0c6e2e82081 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignmentInAncestor_after.py @@ -0,0 +1,8 @@ +class Base(object): + attr: [int] = 0 + + +class MyClass(Base): + def __init___(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignment_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignment_after.py new file mode 100644 index 000000000000..a2fe32a5c168 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInstanceAttributeClassLevelAssignment_after.py @@ -0,0 +1,6 @@ +class MyClass: + attr: [int] = 0 + + def __init___(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignment.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignment.py new file mode 100644 index 000000000000..03f3eb1aad9c --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignment.py @@ -0,0 +1,6 @@ +class MyClass: + attr = 0 + + def __init___(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignmentInAncestor.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignmentInAncestor.py new file mode 100644 index 000000000000..46a47f7f0d59 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignmentInAncestor.py @@ -0,0 +1,8 @@ +class Base(object): + attr = 0 + + +class MyClass(Base): + def __init___(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignmentInAncestor_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignmentInAncestor_after.py new file mode 100644 index 000000000000..ad4184e26fa8 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignmentInAncestor_after.py @@ -0,0 +1,8 @@ +class Base(object): + attr = 0 # type: [int] + + +class MyClass(Base): + def __init___(self): + self.attr = 42 + self.attr \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignment_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignment_after.py new file mode 100644 index 000000000000..a3513d4cfcb5 --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/typeCommentInstanceAttributeClassLevelAssignment_after.py @@ -0,0 +1,6 @@ +class MyClass: + attr = 0 # 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 2c9e5cb237b3..db116eed5ed6 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java @@ -90,6 +90,14 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { doTestAnnotation(); } + public void testAnnotationInstanceAttributeClassLevelAssignment() { + doTestAnnotation(); + } + + public void testAnnotationInstanceAttributeClassLevelAssignmentInAncestor() { + doTestAnnotation(); + } + public void testTypeCommentInstanceAttribute() { doTestTypeComment(); } @@ -102,6 +110,14 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { doTestTypeComment(); } + public void testTypeCommentInstanceAttributeClassLevelAssignment() { + doTestTypeComment(); + } + + public void testTypeCommentInstanceAttributeClassLevelAssignmentInAncestor() { + doTestTypeComment(); + } + public void testTypeCommentLocalSimpleAssignmentTarget() { doTestTypeComment(); }