PY-16412 Start working on the intention to generate variable annotations

2026-09-27 10:03:11 +07:00 · 2018-02-05 21:41:59 +03:00
parent 46cff81cee
commit 0cb6c91134
8 changed files with 144 additions and 0 deletions
@@ -0,0 +1,5 @@
+class C:
+    attr: str
+
+    def __init__(self, attr='spam'):
+        self.attr = attr
@@ -0,0 +1,3 @@
+class C:
+    def __init__(self, attr='spam'):
+        self.attr = attr
@@ -0,0 +1,6 @@
+<html>
+<body>
+<span>This intention adds type hints for variables in PEP 484 compatible format.</span>
+<!-- tooltip end -->
+</body>
+</html>
@@ -340,6 +340,11 @@
      <category>Python</category>
    </intentionAction>

+    <intentionAction>
+      <className>com.jetbrains.python.codeInsight.intentions.PyAnnotateVariableTypeIntention</className>
+      <category>Python</category>
+    </intentionAction>
+
    <testFinder implementation="com.jetbrains.python.codeInsight.testIntegration.PyTestFinder"/>
    <testCreator language="Python" implementationClass="com.jetbrains.python.codeInsight.testIntegration.PyTestCreator"/>

@@ -0,0 +1,92 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.codeInsight.intentions;
+
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.roots.ProjectFileIndex;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider;
+import com.jetbrains.python.documentation.doctest.PyDocstringFile;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.resolve.PyResolveContext;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import one.util.streamex.StreamEx;
+import org.jetbrains.annotations.Nls;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction {
+  @Nls
+  @NotNull
+  @Override
+  public String getFamilyName() {
+    return PyBundle.message("INTN.annotate.types");
+  }
+
+  @Override
+  public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+    if (!(file instanceof PyFile) || file instanceof PyDocstringFile) {
+      return false;
+    }
+    final List<PyTargetExpression> resolved = findSuitableTargetsUnderCaret(project, editor, file);
+    if (resolved.isEmpty() || resolved.size() > 1) {
+      return false;
+    }
+
+    setText(PyBundle.message("INTN.annotate.types"));
+    return true;
+  }
+
+  @NotNull
+  private static List<PyTargetExpression> findSuitableTargetsUnderCaret(@NotNull Project project, Editor editor, PsiFile file) {
+    final PyReferenceOwner elementAtCaret = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()),
+                                                                        PyReferenceExpression.class, PyTargetExpression.class);
+    if (elementAtCaret == null) {
+      return Collections.emptyList();
+    }
+
+    final ProjectFileIndex index = ProjectFileIndex.getInstance(project);
+    final TypeEvalContext typeEvalContext = TypeEvalContext.codeAnalysis(project, file);
+    final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(typeEvalContext);
+    // TODO filter out targets defined in stubs
+    return StreamEx.of(PyUtil.multiResolveTopPriority(elementAtCaret, resolveContext))
+      .select(PyTargetExpression.class)
+      .filter(target -> !index.isInLibraryClasses(target.getContainingFile().getVirtualFile()))
+      .filter(target -> !isAnnotated(target, typeEvalContext))
+      .toList();
+  }
+
+  private static boolean isAnnotated(@NotNull PyTargetExpression target, @NotNull TypeEvalContext context) {
+    // TODO filter out fields explicitly annotated as Any
+    return new PyTypingTypeProvider().getReferenceType(target, context, null) != null;
+  }
+
+  @Override
+  public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+    final List<PyTargetExpression> targets = findSuitableTargetsUnderCaret(project, editor, file);
+    assert targets.size() == 1;
+    final PyTargetExpression annotationTarget = targets.get(0);
+    if (LanguageLevel.forElement(annotationTarget).isAtLeast(LanguageLevel.PYTHON36)) {
+      insertVariableAnnotation(annotationTarget);
+    }
+    else {
+      insertVariableTypeComment(annotationTarget);
+    }
+  }
+
+  private void insertVariableAnnotation(@NotNull PyTargetExpression target) {
+
+  }
+
+  private void insertVariableTypeComment(@NotNull PyTargetExpression target) {
+
+  }
+}
@@ -0,0 +1,3 @@
+def func():
+    var = 'spam'
+    v<caret>ar
@@ -0,0 +1,3 @@
+def func():
+    var: str = 'spam'
+    var
@@ -0,0 +1,27 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.intentions;
+
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.psi.LanguageLevel;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase {
+  public void testAnnotationLocalAssignmentTarget() {
+    doTestAnnotation();
+  }
+
+  private void doTestAnnotation() {
+    doTest(LanguageLevel.PYTHON36);
+  }
+
+  private void doTestTypeComment() {
+    doTest(LanguageLevel.PYTHON27);
+  }
+
+  private void doTest(@NotNull LanguageLevel languageLevel) {
+    doTest(PyBundle.message("INTN.annotate.types"), languageLevel);
+  }
+}