PY-16412 Extract common logic for generation of variable annotations

2026-09-27 10:03:11 +07:00 · 2018-02-05 21:41:59 +03:00
parent f2a2a4a070
commit 54574eac34
16 changed files with 111 additions and 20 deletions
@@ -9,9 +9,11 @@ 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.PythonDocumentationProvider;
 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.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.Nls;
@@ -74,7 +76,7 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction {
    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)) {
+    if (preferSyntacticAnnotation(annotationTarget)) {
      insertVariableAnnotation(annotationTarget);
    }
    else {
@@ -82,8 +84,15 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction {
    }
  }

-  private void insertVariableAnnotation(@NotNull PyTargetExpression target) {
+  private static boolean preferSyntacticAnnotation(@NotNull PyTargetExpression annotationTarget) {
+    return LanguageLevel.forElement(annotationTarget).isAtLeast(LanguageLevel.PYTHON36);
+  }

+  private static void insertVariableAnnotation(@NotNull PyTargetExpression target) {
+    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);
  }

  private void insertVariableTypeComment(@NotNull PyTargetExpression target) {
@@ -35,7 +35,6 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;

 import java.util.ArrayList;
-import java.util.List;
 import java.util.Map;

 public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseIntentionAction {
@@ -48,24 +47,11 @@ public class PyConvertTypeCommentToVariableAnnotationIntention extends PyBaseInt
        if (typeComment.getParent() instanceof PyAssignmentStatement && map.size() == 1) {
          final PyTargetExpression target = ContainerUtil.getFirstItem(map.keySet());
          assert target != null;
-          PyUtil.updateDocumentUnblockedAndCommitted(target, document -> {
-            document.insertString(target.getTextRange().getEndOffset(), ": " + map.get(target));
-          });
+          PyTypeHintGenerationUtil.insertVariableAnnotation(target, map.get(target));
        }
        else {
-          final PyStatement statement = PsiTreeUtil.getParentOfType(typeComment, PyStatement.class);
-          assert statement != null;
-
-          final PyElementGenerator generator = PyElementGenerator.getInstance(project);
-          final List<Map.Entry<PyTargetExpression, String>> entries = new ArrayList<>(map.entrySet());
-
-          for (Map.Entry<PyTargetExpression, String> entry : entries) {
-            final PyTargetExpression target = entry.getKey();
-            final String annotation = entry.getValue();
-            final PyTypeDeclarationStatement declaration = generator.createFromText(LanguageLevel.PYTHON36,
-                                                                                    PyTypeDeclarationStatement.class,
-                                                                                    target.getText() + ": " + annotation);
-            statement.getParent().addBefore(declaration, statement);
+          for (Map.Entry<PyTargetExpression, String> entry : new ArrayList<>(map.entrySet())) {
+            PyTypeHintGenerationUtil.insertVariableAnnotation(entry.getKey(), entry.getValue());
          }
        }

@@ -0,0 +1,41 @@
+// 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.psi.util.PsiTreeUtil;
+import com.jetbrains.python.psi.*;
+import org.jetbrains.annotations.NotNull;
+
+import static com.jetbrains.python.psi.PyUtil.as;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PyTypeHintGenerationUtil {
+  private PyTypeHintGenerationUtil() {}
+
+  public static void insertVariableAnnotation(@NotNull PyTargetExpression target, @NotNull String annotation) {
+    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);
+    }
+
+    if (canUseInlineAnnotation(target)) {
+      PyUtil.updateDocumentUnblockedAndCommitted(target, document -> {
+        document.insertString(target.getTextRange().getEndOffset(), ": " + annotation);
+      });
+    }
+    else {
+      final PyElementGenerator generator = PyElementGenerator.getInstance(target.getProject());
+      final String declarationText = target.getText() + ": " + annotation;
+      final PyTypeDeclarationStatement declaration = generator.createFromText(langLevel, PyTypeDeclarationStatement.class, declarationText);
+      final PyStatement statement = PsiTreeUtil.getParentOfType(target, PyStatement.class);
+      assert statement != null;
+      statement.getParent().addBefore(declaration, statement);
+    }
+  }
+
+  private static boolean canUseInlineAnnotation(@NotNull PyTargetExpression target) {
+    final PyAssignmentStatement assignment = as(target.getParent(), PyAssignmentStatement.class);
+    return assignment != null && assignment.getRawTargets().length == 1 && assignment.getLeftHandSideExpression() == target;
+  }
+}
@@ -0,0 +1,3 @@
+def func():
+    var = another = 'spam'
+    v<caret>ar
@@ -0,0 +1,4 @@
+def func():
+    var: str
+    var = another = 'spam'
+    var
@@ -0,0 +1,3 @@
+def func():
+    for var in 'spam':
+        v<caret>ar
@@ -0,0 +1,4 @@
+def func():
+    var: str
+    for var in 'spam':
+        var
@@ -0,0 +1,3 @@
+def func():
+    (var) = 'spam'
+    v<caret>ar
@@ -0,0 +1,4 @@
+def func():
+    var: Any
+    (var) = 'spam'
+    var
@@ -0,0 +1,3 @@
+def func():
+    var, _ = 'spam', 42
+    v<caret>ar
@@ -0,0 +1,4 @@
+def func():
+    var: str
+    var, _ = 'spam', 42
+    var
@@ -0,0 +1,3 @@
+def func():
+    with open('file.txt') as var:
+        v<caret>ar
@@ -0,0 +1,4 @@
+def func():
+    var: TextIO
+    with open('file.txt') as var:
+        var
@@ -9,7 +9,27 @@ import org.jetbrains.annotations.NotNull;
 * @author Mikhail Golubev
 */
 public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase {
-  public void testAnnotationLocalAssignmentTarget() {
+  public void testAnnotationLocalSimpleAssignmentTarget() {
+    doTestAnnotation();
+  }
+
+  public void testAnnotationLocalSimpleAssignmentTargetInParentheses() {
+    doTestAnnotation();
+  }
+
+  public void testAnnotationLocalUnpackedAssignmentTarget() {
+    doTestAnnotation();
+  }
+
+  public void testAnnotationLocalChainedAssignmentTarget() {
+    doTestAnnotation();
+  }
+
+  public void testAnnotationLocalForTarget() {
+    doTestAnnotation();
+  }
+
+  public void testAnnotationLocalWithTarget() {
    doTestAnnotation();
  }