Intention to add PEP484 Py3.5-style type annotations to a function

2026-09-27 10:03:11 +07:00 · 2016-02-22 22:38:42 +01:00
parent 80acb80e01
commit 2697fdc568
29 changed files with 385 additions and 43 deletions
@@ -0,0 +1,2 @@
+def foo(var: object) -> object:
+    pass
@@ -0,0 +1,2 @@
+def foo(var):
+    pass
@@ -0,0 +1,9 @@
+<html>
+<body>
+<span style="font-family: verdana,sans-serif;">
+  This intention used to specify annotations for parameters and return type of a function.
+  <p/>
+  If there is type information collected in run-time, it is used to set the default values of types.
+</span>
+</body>
+</html>
@@ -185,6 +185,11 @@
      <category>Python</category>
    </intentionAction>

+    <intentionAction>
+      <className>com.jetbrains.python.codeInsight.intentions.PyAnnotateTypesIntention</className>
+      <category>Python</category>
+    </intentionAction>
+
    <intentionAction>
      <className>com.jetbrains.python.codeInsight.intentions.TypeAssertionIntention</className>
      <category>Python</category>
@@ -258,6 +258,9 @@ INTN.add.parameters.to.docstring=Add parameters to docstring
 INTN.specify.type.in.annotation=Specify type for reference using annotation
 INTN.specify.return.type.in.annotation=Specify return type using annotation

+#PyAnnotateTypesIntention
+INTN.annotate.types=Annotate types
+
 #TypeAssertionIntention
 INTN.insert.assertion=Insert type assertion

@@ -0,0 +1,135 @@
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.codeInsight.intentions;
+
+import com.google.common.base.Function;
+import com.intellij.codeInsight.CodeInsightUtilCore;
+import com.intellij.codeInsight.intention.IntentionAction;
+import com.intellij.codeInsight.template.*;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.fileEditor.FileEditorManager;
+import com.intellij.openapi.fileEditor.OpenFileDescriptor;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.documentation.doctest.PyDocstringFile;
+import com.jetbrains.python.psi.*;
+import org.jetbrains.annotations.NotNull;
+
+import static com.jetbrains.python.codeInsight.intentions.SpecifyTypeInPy3AnnotationsIntention.annotateParameter;
+import static com.jetbrains.python.codeInsight.intentions.SpecifyTypeInPy3AnnotationsIntention.annotateReturnType;
+import static com.jetbrains.python.codeInsight.intentions.TypeIntention.getCallable;
+import static com.jetbrains.python.codeInsight.intentions.TypeIntention.resolvesToFunction;
+
+/**
+ * @author traff
+ */
+public class PyAnnotateTypesIntention implements IntentionAction {
+  private String myText = PyBundle.message("INTN.annotate.types");
+
+  @NotNull
+  public String getText() {
+    return myText;
+  }
+
+  @NotNull
+  public String getFamilyName() {
+    return PyBundle.message("INTN.annotate.types");
+  }
+
+  public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+    if (!(file instanceof PyFile) || file instanceof PyDocstringFile) return false;
+
+    updateText();
+
+    final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset());
+    if (elementAt == null) return false;
+
+    if (resolvesToFunction(elementAt, new Function<PyFunction, Boolean>() {
+      @Override
+      public Boolean apply(PyFunction input) {
+        return true;
+      }
+    })) {
+      updateText();
+      return true;
+    }
+    return false;
+  }
+
+  public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+    final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset());
+    final PyCallable callable = getCallable(elementAt);
+
+
+    final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(callable);
+
+    PyExpression returnType = annotateReturnType(project, editor.getDocument(), elementAt, false);
+
+    if (returnType != null) {
+      builder.replaceElement(returnType, returnType.getText());
+    }
+
+    if (callable instanceof PyFunction) {
+      PyFunction function = (PyFunction) callable;
+      PyParameter[] params = function.getParameterList().getParameters();
+
+      for (int i = params.length - 1; i >= 0; i--) {
+        if (params[i] instanceof PyNamedParameter) {
+          params[i] = annotateParameter(project, editor, (PyNamedParameter)params[i], false);
+        }
+      }
+
+
+      for (int i = params.length - 1; i >= 0; i--) {
+        if (params[i] instanceof PyNamedParameter) {
+          params[i] = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(params[i]);
+          PyAnnotation annotation = ((PyNamedParameter)params[i]).getAnnotation();
+          if (annotation != null) {
+            PyExpression annotationValue = annotation.getValue();
+            builder.replaceElement(annotationValue, annotationValue.getText());
+          }
+        }
+      }
+    }
+    final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
+
+    int offset = callable.getTextRange().getStartOffset();
+
+    final OpenFileDescriptor descriptor = new OpenFileDescriptor(
+      project,
+      callable.getContainingFile().getVirtualFile(),
+      offset
+    );
+    final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
+    if (targetEditor != null) {
+      targetEditor.getCaretModel().moveToOffset(offset);
+      TemplateManager.getInstance(project).startTemplate(targetEditor, template);
+    }
+  }
+
+  @Override
+  public boolean startInWriteAction() {
+    return true;
+  }
+
+
+  protected void updateText() {
+    myText = PyBundle.message("INTN.annotate.types");
+  }
+}
@@ -73,14 +73,17 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
    final PyNamedParameter parameter = getParameter(problemElement, resolved);

    if (parameter != null) {
-      annotateParameter(project, editor, parameter);
+      annotateParameter(project, editor, parameter, true);
    }
    else {
-      annotateReturnType(project, editor.getDocument(), elementAt);
+      annotateReturnType(project, editor.getDocument(), elementAt, true);
    }
  }

-  private static void annotateParameter(Project project, Editor editor, @NotNull PyNamedParameter parameter) {
+  static PyNamedParameter annotateParameter(Project project,
+                                                   Editor editor,
+                                                   @NotNull PyNamedParameter parameter,
+                                                   boolean createTemplate) {
    final PyExpression defaultParamValue = parameter.getDefaultValue();

    final String paramName = StringUtil.notNullize(parameter.getName());
@@ -106,7 +109,7 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
    parameter = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(parameter);
    editor.getCaretModel().moveToOffset(parameter.getTextOffset());
    final PyAnnotation annotation = parameter.getAnnotation();
-    if (annotation != null) {
+    if (annotation != null && createTemplate) {
      final PyExpression annotationValue = annotation.getValue();

      final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(parameter);
@@ -117,42 +120,42 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
      final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
      TemplateManager.getInstance(project).startTemplate(editor, template);
    }
+
+    return parameter;
  }

-  private void annotateReturnType(Project project, Document document, PsiElement resolved) {
+  public static PyExpression annotateReturnType(Project project, Document document, PsiElement resolved, boolean createTemplate) {
    PyCallable callable = getCallable(resolved);

    String returnType = PyNames.OBJECT;

    if (callable instanceof PyFunction) {
-      PyFunction function = (PyFunction) callable;
+      PyFunction function = (PyFunction)callable;
      final PySignature signature = PySignatureCacheManager.getInstance(project).findSignature(
        function);
      if (signature != null) {
        returnType = ObjectUtils.chooseNotNull(signature.getReturnTypeQualifiedName(), returnType);
      }
-    }

-
-    if (callable instanceof PyFunction) {
      final String annotationText = " -> " + returnType;

      final PsiElement prevElem = PyPsiUtils.getPrevNonCommentSibling(((PyFunction)callable).getStatementList(), true);
      assert prevElem != null;

      final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
+      Document documentWithCollable = manager.getDocument(callable.getContainingFile());
      try {
        final TextRange range = prevElem.getTextRange();
-        manager.doPostponedOperationsAndUnblockDocument(document);
+        manager.doPostponedOperationsAndUnblockDocument(documentWithCollable);
        if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) {
-          document.insertString(range.getStartOffset(), annotationText);
+          documentWithCollable.insertString(range.getStartOffset(), annotationText);
        }
        else {
-          document.insertString(range.getEndOffset(), annotationText + ":");
+          documentWithCollable.insertString(range.getEndOffset(), annotationText + ":");
        }
      }
      finally {
-        manager.commitDocument(document);
+        manager.commitDocument(documentWithCollable);
      }


@@ -161,22 +164,28 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
      assert annotation != null;
      final PyExpression annotationValue = annotation.getValue();
      assert annotationValue != null : "Generated function must have annotation";
-      final int offset = annotationValue.getTextOffset();

-      final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(annotationValue);
-      builder.replaceRange(TextRange.create(0, returnType.length()), returnType);
-      final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
-      final OpenFileDescriptor descriptor = new OpenFileDescriptor(
-        project,
-        callable.getContainingFile().getVirtualFile(),
-        offset
-      );
-      final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
-      if (targetEditor != null) {
-        targetEditor.getCaretModel().moveToOffset(offset);
-        TemplateManager.getInstance(project).startTemplate(targetEditor, template);
+      if (createTemplate) {
+        final int offset = annotationValue.getTextOffset();
+
+        final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(annotationValue);
+        builder.replaceRange(TextRange.create(0, returnType.length()), returnType);
+        final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
+        final OpenFileDescriptor descriptor = new OpenFileDescriptor(
+          project,
+          callable.getContainingFile().getVirtualFile(),
+          offset
+        );
+        final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
+        if (targetEditor != null) {
+          targetEditor.getCaretModel().moveToOffset(offset);
+          TemplateManager.getInstance(project).startTemplate(targetEditor, template);
+        }
      }
+      return annotationValue;
    }
+
+    return null;
  }

  @Override
@@ -15,6 +15,7 @@
 */
 package com.jetbrains.python.codeInsight.intentions;

+import com.google.common.base.Function;
 import com.intellij.codeInsight.intention.IntentionAction;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.editor.Editor;
@@ -80,7 +81,7 @@ public abstract class TypeIntention implements IntentionAction {
  }

  @Nullable
-  protected static PyExpression getProblemElement(@Nullable PsiElement elementAt) {
+  public static PyExpression getProblemElement(@Nullable PsiElement elementAt) {
    PyExpression problemElement = PsiTreeUtil.getParentOfType(elementAt, PyNamedParameter.class, PyReferenceExpression.class);
    if (problemElement == null) return null;
    if (problemElement instanceof PyQualifiedExpression) {
@@ -108,13 +109,22 @@ public abstract class TypeIntention implements IntentionAction {
  }

  private boolean isAvailableForReturn(@NotNull final PsiElement elementAt) {
+    return resolvesToFunction(elementAt, new Function<PyFunction, Boolean>() {
+      @Override
+      public Boolean apply(PyFunction input) {
+        return !isReturnTypeDefined(input);
+      }
+    });
+  }
+
+  static boolean resolvesToFunction(@NotNull PsiElement elementAt, Function<PyFunction, Boolean> isAvailableForFunction) {
    final PyFunction parentFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
    if (parentFunction != null) {
      final ASTNode nameNode = parentFunction.getNameNode();
      if (nameNode != null) {
        final PsiElement prev = elementAt.getContainingFile().findElementAt(elementAt.getTextOffset()-1);
        if (nameNode.getPsi() == elementAt || nameNode.getPsi() == prev) {
-          return !isReturnTypeDefined(parentFunction);
+          return  isAvailableForFunction.apply(parentFunction);
        }
      }
    }
@@ -126,18 +136,19 @@ public abstract class TypeIntention implements IntentionAction {
    final PsiReference reference = callee.getReference();
    if (reference instanceof PsiPolyVariantReference) {
      final ResolveResult[] results = ((PsiPolyVariantReference)reference).multiResolve(false);
-      if (results.length == 1) {
-        final PsiElement result = results[0].getElement();
-        if (!(result instanceof PyFunction)) return false;
-        final PsiFile psiFile = result.getContainingFile();
-        if (psiFile == null) return false;
-        final VirtualFile virtualFile = psiFile.getVirtualFile();
-        if (virtualFile != null) {
-          if (ProjectRootManager.getInstance(psiFile.getProject()).getFileIndex().isInLibraryClasses(virtualFile)) {
-            return false;
+      for (int i = 0; i<results.length; i++) {
+        if (results[i].getElement() instanceof PyFunction) {
+          final PsiElement result = results[i].getElement();
+          final PsiFile psiFile = result.getContainingFile();
+          if (psiFile == null) return false;
+          final VirtualFile virtualFile = psiFile.getVirtualFile();
+          if (virtualFile != null) {
+            if (ProjectRootManager.getInstance(psiFile.getProject()).getFileIndex().isInLibraryClasses(virtualFile)) {
+              return false;
+            }
          }
+          return isAvailableForFunction.apply((PyFunction)result);
        }
-        return !isReturnTypeDefined((PyFunction)result);
      }
    }
    return false;
@@ -148,7 +159,7 @@ public abstract class TypeIntention implements IntentionAction {
  }

  @Nullable
-  protected static PyCallExpression getCallExpression(PsiElement elementAt) {
+  static PyCallExpression getCallExpression(PsiElement elementAt) {
    final PyExpression problemElement = getProblemElement(elementAt);
    if (problemElement != null) {
      PsiReference reference = problemElement.getReference();
@@ -172,7 +183,7 @@ public abstract class TypeIntention implements IntentionAction {
  }

  @Nullable
-  protected PyCallable getCallable(PsiElement elementAt) {
+  static PyCallable getCallable(PsiElement elementAt) {
    PyCallExpression callExpression = getCallExpression(elementAt);

    if (callExpression != null && elementAt != null) {
@@ -182,7 +193,7 @@ public abstract class TypeIntention implements IntentionAction {
    return PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
  }

-  protected PyResolveContext getResolveContext(@NotNull PsiElement origin) {
+  protected static PyResolveContext getResolveContext(@NotNull PsiElement origin) {
    return PyResolveContext.defaultContext().withTypeEvalContext(TypeEvalContext.codeAnalysis(origin.getProject(), origin.getContainingFile()));
  }

@@ -0,0 +1,2 @@
+def fo<caret>o(x, y):
+    pass
@@ -0,0 +1,2 @@
+def foo(x: object, y: object) -> object:
+    pass
@@ -0,0 +1,3 @@
+from foo_decl import foo
+
+fo<caret>o(1, 1)
@@ -0,0 +1,3 @@
+from foo_decl import foo
+
+foo(1, 1)
@@ -0,0 +1,5 @@
+def foo(x, y):
+    pass
+
+
+fo<caret>o(1, 1)
@@ -0,0 +1,5 @@
+def foo(x: object, y: object) -> object:
+    pass
+
+
+foo(1, 1)
@@ -0,0 +1,2 @@
+def foo(x, y):
+    pass
@@ -0,0 +1,2 @@
+def foo(x: object, y: object) -> object:
+    pass
@@ -0,0 +1,2 @@
+def fo<caret>o(x, y):
+    pass
@@ -0,0 +1,2 @@
+def foo(x, y) -> object:
+    pass
@@ -0,0 +1,3 @@
+from foo_decl import foo
+
+fo<caret>o(1, 1)
@@ -0,0 +1,3 @@
+from foo_decl import foo
+
+foo(1, 1)
@@ -0,0 +1,5 @@
+def foo(x, y):
+    pass
+
+
+fo<caret>o(1, 1)
@@ -0,0 +1,5 @@
+def foo(x, y) -> object:
+    pass
+
+
+foo(1, 1)
@@ -0,0 +1,3 @@
+def foo(var):
+    print(va<caret>r)
+    pass
@@ -0,0 +1,3 @@
+def foo(var: object):
+    print(var)
+    pass
@@ -0,0 +1,2 @@
+def foo(x, y):
+    pass
@@ -0,0 +1,2 @@
+def foo(x, y) -> object:
+    pass
@@ -0,0 +1,48 @@
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.intentions;
+
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.psi.LanguageLevel;
+import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
+
+/**
+ * @author traff
+ */
+public class PyAnnotateTypesIntentionTest extends PyIntentionTestCase {
+  public void testCaretOnDefinition() {
+    doTest();
+  }
+
+  public void testCaretOnInvocation() {
+    doTest();
+  }
+
+  public void testCaretOnImportedInvocation() {
+    PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), LanguageLevel.PYTHON30);
+    try {
+      doIntentionTest(PyBundle.message("INTN.annotate.types"), getTestName(true) + ".py", "foo_decl.py");
+      myFixture.checkResultByFile("foo_decl.py", "foo_decl_after.py", false);
+    }
+    finally {
+      PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null);
+    }
+  }
+
+  private void doTest() {
+    doTest(PyBundle.message("INTN.annotate.types"), LanguageLevel.PYTHON30);
+  }
+}
@@ -42,9 +42,13 @@ public abstract class PyIntentionTestCase extends PyTestCase {
    }
  }

-  protected void doIntentionTest(final String hint) {
+  protected void doIntentionTest(final String hint, String ... files) {
    final String testFileName = getTestName(true);
-    myFixture.configureByFile(testFileName + ".py");
+    if (files.length>0) {
+      myFixture.configureByFiles(files);
+    } else {
+      myFixture.configureByFile(testFileName + ".py");
+    }
    final IntentionAction intentionAction = myFixture.findSingleIntention(hint);
    assertNotNull(intentionAction);
    myFixture.launchAction(intentionAction);
@@ -0,0 +1,60 @@
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.intentions;
+
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.psi.LanguageLevel;
+import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
+
+/**
+ * @author traff
+ */
+public class SpecifyTypeInPy3AnnotationsIntentionTest extends PyIntentionTestCase {
+  public void testCaretOnDefinition() {
+    doTestReturnType();
+  }
+
+
+
+  public void testCaretOnInvocation() {
+    doTestReturnType();
+  }
+
+  public void testCaretOnImportedInvocation() {
+    PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), LanguageLevel.PYTHON30);
+    try {
+      doIntentionTest(PyBundle.message("INTN.specify.return.type.in.annotation"), getTestName(true) + ".py", "foo_decl.py");
+      myFixture.checkResultByFile("foo_decl.py", "foo_decl_after.py", false);
+    }
+    finally {
+      PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null);
+    }
+  }
+
+  public void testCaretOnParamUsage() {
+    doTestParam();
+  }
+
+
+  private void doTestReturnType() {
+    doTest(PyBundle.message("INTN.specify.return.type.in.annotation"), LanguageLevel.PYTHON30);
+  }
+
+
+  private void doTestParam() {
+    doTest(PyBundle.message("INTN.specify.type.in.annotation"), LanguageLevel.PYTHON30);
+  }
+}