PY-16412 Support type hints for instance attributes using annotations

2026-09-27 10:03:11 +07:00 · 2018-02-05 21:41:59 +03:00
parent af24ba5f20
commit 1d71f182c1
13 changed files with 174 additions and 45 deletions
@@ -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<? extends RatedResolveResult> 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<? extends RatedResolveResult> 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) {
@@ -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);
    }
  }

@@ -0,0 +1,4 @@
+class MyClass:
+    def __init__(self):
+        self.attr = 42
+        self.at<caret>tr
@@ -0,0 +1,6 @@
+class MyClass:
+    """Docstring."""
+
+    def __init__(self):
+        self.attr = 42
+        self.at<caret>tr
@@ -0,0 +1,7 @@
+class MyClass:
+    """Docstring."""
+    attr: [int]
+
+    def __init__(self):
+        self.attr = 42
+        self.attr
@@ -0,0 +1,6 @@
+class MyClass:
+    attr: [int]
+
+    def __init__(self):
+        self.attr = 42
+        self.attr
@@ -0,0 +1,6 @@
+class MyClass:
+    attr: int
+
+    def __init___(self, x):
+        self.attr = x
+        self.at<caret>tr
@@ -0,0 +1,6 @@
+class MyClass:
+    attr = ...  # type: int
+
+    def __init___(self, x):
+        self.attr = x
+        self.at<caret>tr
@@ -0,0 +1,4 @@
+class MyClass:
+    def __init__(self):
+        self.attr = 42
+        self.at<caret>tr
@@ -0,0 +1,6 @@
+class MyClass:
+    """Docstring."""
+
+    def __init__(self):
+        self.attr = 42
+        self.at<caret>tr
@@ -0,0 +1,7 @@
+class MyClass:
+    """Docstring."""
+    attr = None  # type: [int]
+
+    def __init__(self):
+        self.attr = 42
+        self.at<caret>tr
@@ -0,0 +1,6 @@
+class MyClass:
+    attr = None  # type: [int]
+
+    def __init__(self):
+        self.attr = 42
+        self.attr
@@ -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();
  }