fixed PY-7354 "Add return type annotation" quickfix on function to add :rtype: to docstring

python/src/com/jetbrains/python/PyBundle.properties

@@ -184,9 +184,11 @@ INTN.doc.string.stub=Insert documentation string stub
 
 #SpecifyTypeInDocstringIntention
 INTN.specify.type=Specify type for reference in docstring
+INTN.specify.return.type=Specify return type in docstring
 
 #SpecifyTypeInPy3AnnotationsIntention
 INTN.specify.type.in.annotation=Specify type for reference using annotation
+INTN.specify.returt.type.in.annotation=Specify return type using annotation
 
 #TypeAssertionIntention
 INTN.insert.assertion=Insert type assertion

python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java

@@ -3,17 +3,20 @@ package com.jetbrains.python.codeInsight.intentions;
 import com.intellij.codeInsight.CodeInsightUtilBase;
 import com.intellij.codeInsight.intention.IntentionAction;
 import com.intellij.codeInsight.template.*;
+import com.intellij.lang.ASTNode;
 import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiReference;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyNames;
+import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.documentation.PyDocumentationSettings;
 import com.jetbrains.python.documentation.PythonDocumentationProvider;
 import com.jetbrains.python.psi.*;
@@ -28,13 +31,13 @@ import org.jetbrains.annotations.NotNull;
  * Helps to specify type
  */
 public class SpecifyTypeInDocstringIntention implements IntentionAction {
-
+  private String myText = PyBundle.message("INTN.specify.type");
   public SpecifyTypeInDocstringIntention() {
   }
 
   @NotNull
   public String getText() {
-    return PyBundle.message("INTN.specify.type");
+    return myText;
   }
 
   @NotNull
@@ -43,6 +46,19 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction {
   }
 
   public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+    PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1);
+    if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER))
+      elementAt = file.findElementAt(editor.getCaretModel().getOffset());
+
+    PyFunction parentFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
+    if (parentFunction != null) {
+      final ASTNode nameNode = parentFunction.getNameNode();
+      if (nameNode != null && nameNode.getPsi() == elementAt) {
+        myText = PyBundle.message("INTN.specify.return.type");
+        return true;
+      }
+    }
+
     PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class);
 
     if (problemElement == null) return false;
@@ -82,10 +98,33 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction {
   }
 
   public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+    PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1);
+    if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER))
+      elementAt = file.findElementAt(editor.getCaretModel().getOffset());
+
+    String type = "type";
+    String name = "";
+    PyFunction pyFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
+    if (pyFunction != null) {
+      final ASTNode nameNode = pyFunction.getNameNode();
+      if (nameNode != null && nameNode.getPsi() == elementAt) {
+        type = "rtype";
+      }
+    }
+
+    PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(project);
+    String prefix = ":";
+    if (documentationSettings.isEpydocFormat(file)) {
+      prefix = "@";
+    }
+
+    PsiReference reference = null;
     PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class);
+    PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
     if (problemElement != null) {
-      String name = problemElement.getName();
-      PsiReference reference = problemElement.getReference();
+      name = problemElement.getName();
+
+      reference = problemElement.getReference();
       if (problemElement instanceof PyQualifiedExpression) {
         final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier();
         if (qualifier != null) {
@@ -93,74 +132,68 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction {
           name = qualifier.getText();
         }
       }
-      PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(project);
-      String prefix = ":";
-      if (documentationSettings.isEpydocFormat(file)) {
-        prefix = "@";
-      }
-      PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
+      pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class);
+    }
 
-      PyFunction pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class);
+    final ASTNode nameNode = pyFunction.getNameNode();
+    if ((pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter)) ||
+      elementAt == nameNode.getPsi()) {
+      PyStringLiteralExpression docStringExpression = pyFunction.getDocStringExpression();
+      int startOffset;
+      int endOffset;
+      final Pair<String, Integer> replacementToOffset =
+          PythonDocumentationProvider.addParamToDocstring(pyFunction, type, name, prefix);
 
-      if (pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter)) {
-        PyStringLiteralExpression docStringExpression = pyFunction.getDocStringExpression();
-        int startOffset;
-        int endOffset;
-        final Pair<String, Integer> replacementToOffset =
-            PythonDocumentationProvider.addParamToDocstring(pyFunction, "type", name, prefix);
-
-        if (docStringExpression != null) {
-          final String typePattern = "type " + name + ":";
-          final int index = docStringExpression.getText().indexOf(typePattern);
-          if (index == -1) {
-            PyExpression str = elementGenerator.createDocstring(replacementToOffset.getFirst()).getExpression();
-            docStringExpression.replace(str);
-            startOffset = replacementToOffset.getSecond();
-            endOffset = startOffset;
-          }
-          else {
-            startOffset = index + typePattern.length() + 1;
-            endOffset = docStringExpression.getText().indexOf("\n", startOffset);
-            if (endOffset == -1) endOffset = startOffset;
-          }
-          pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction);
-          docStringExpression = pyFunction.getDocStringExpression();
-        }
-        else {
-
-          final PyStatementList list = pyFunction.getStatementList();
-          final Document document = editor.getDocument();
+      if (docStringExpression != null) {
+        final String typePattern = type + " " + name + ":";
+        final int index = docStringExpression.getText().indexOf(typePattern);
+        if (index == -1) {
+          PyExpression str = elementGenerator.createDocstring(replacementToOffset.getFirst()).getExpression();
+          docStringExpression.replace(str);
           startOffset = replacementToOffset.getSecond();
-
-          if (list != null && list.getStatements().length != 0) {
-            if (document.getLineNumber(list.getTextOffset()) == document.getLineNumber(pyFunction.getTextOffset())) {
-              PyFunction func = elementGenerator.createFromText(LanguageLevel.forElement(pyFunction),
-                                        PyFunction.class, "def " + pyFunction.getName() + pyFunction.getParameterList().getText()
-                                        +":\n\t"+replacementToOffset.getFirst() + "\n\t" + list.getText());
-
-              pyFunction = (PyFunction)pyFunction.replace(func);
-              startOffset = replacementToOffset.getSecond() + 2;
-            }
-            else {
-              PyExpressionStatement str = elementGenerator.createDocstring(replacementToOffset.getFirst());
-              list.addBefore(str, list.getStatements()[0]);
-            }
-          }
-
-          pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction);
-          docStringExpression = pyFunction.getDocStringExpression();
-
           endOffset = startOffset;
         }
-        assert docStringExpression != null;
-        int textOffSet = docStringExpression.getTextOffset();
-        editor.getCaretModel().moveToOffset(textOffSet);
-
-        final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(docStringExpression);
-        builder.replaceRange(TextRange.create(startOffset, endOffset), PyNames.OBJECT);
-        Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
-        TemplateManager.getInstance(project).startTemplate(editor, template);
+        else {
+          startOffset = index + typePattern.length() + 1;
+          endOffset = docStringExpression.getText().indexOf("\n", startOffset);
+          if (endOffset == -1) endOffset = startOffset;
+        }
+        pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction);
+        docStringExpression = pyFunction.getDocStringExpression();
       }
+      else {
+        final PyStatementList list = pyFunction.getStatementList();
+        final Document document = editor.getDocument();
+        startOffset = replacementToOffset.getSecond();
+
+        if (list != null && list.getStatements().length != 0) {
+          if (document.getLineNumber(list.getTextOffset()) == document.getLineNumber(pyFunction.getTextOffset())) {
+            PyFunction func = elementGenerator.createFromText(LanguageLevel.forElement(pyFunction),
+                                      PyFunction.class, "def " + pyFunction.getName() + pyFunction.getParameterList().getText()
+                                      +":\n\t"+replacementToOffset.getFirst() + "\n\t" + list.getText());
+
+            pyFunction = (PyFunction)pyFunction.replace(func);
+            startOffset = replacementToOffset.getSecond() + 2;
+          }
+          else {
+            PyExpressionStatement str = elementGenerator.createDocstring(replacementToOffset.getFirst());
+            list.addBefore(str, list.getStatements()[0]);
+          }
+        }
+
+        pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction);
+        docStringExpression = pyFunction.getDocStringExpression();
+
+        endOffset = startOffset;
+      }
+      assert docStringExpression != null;
+      int textOffSet = docStringExpression.getTextOffset();
+      editor.getCaretModel().moveToOffset(textOffSet);
+
+      final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(docStringExpression);
+      builder.replaceRange(TextRange.create(startOffset, endOffset), PyNames.OBJECT);
+      Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
+      TemplateManager.getInstance(project).startTemplate(editor, template);
     }
   }
 

python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java

@@ -3,6 +3,7 @@ package com.jetbrains.python.codeInsight.intentions;
 import com.intellij.codeInsight.CodeInsightUtilBase;
 import com.intellij.codeInsight.intention.IntentionAction;
 import com.intellij.codeInsight.template.*;
+import com.intellij.lang.ASTNode;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.fileEditor.OpenFileDescriptor;
@@ -15,6 +16,7 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyNames;
+import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.types.PyReturnTypeReference;
@@ -28,13 +30,13 @@ import org.jetbrains.annotations.NotNull;
  * Helps to specify type  in annotations in python3
  */
 public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction {
-
+  private String myText = PyBundle.message("INTN.specify.type.in.annotation");
   public SpecifyTypeInPy3AnnotationsIntention() {
   }
 
   @NotNull
   public String getText() {
-    return PyBundle.message("INTN.specify.type.in.annotation");
+    return myText;
   }
 
   @NotNull
@@ -44,6 +46,18 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction {
 
   public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
     if (!LanguageLevel.forElement(file).isPy3K()) return false;
+    PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1);
+    if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER))
+      elementAt = file.findElementAt(editor.getCaretModel().getOffset());
+
+    PyFunction parentFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
+    if (parentFunction != null) {
+      final ASTNode nameNode = parentFunction.getNameNode();
+      if (nameNode != null && nameNode.getPsi() == elementAt) {
+        myText = PyBundle.message("INTN.specify.returt.type.in.annotation");
+        return true;
+      }
+    }
 
     PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class);
 
@@ -96,9 +110,12 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction {
 
   public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
     PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class);
+    PyParameter parameter = null;
+    PsiReference reference = null;
+    String name = null;
     if (problemElement != null) {
-      String name = problemElement.getName();
-      PsiReference reference = problemElement.getReference();
+      name = problemElement.getName();
+      reference = problemElement.getReference();
       if (problemElement instanceof PyQualifiedExpression) {
         final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier();
         if (qualifier != null) {
@@ -106,70 +123,76 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction {
           name = qualifier.getText();
         }
       }
-      PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
+    }
+    final PsiElement resolvedReference = reference != null? reference.resolve() : null;
+    if (problemElement instanceof PyParameter)
+      parameter = (PyParameter)problemElement;
+    else {
+      if (resolvedReference instanceof PyParameter) {
+        parameter = (PyParameter)resolvedReference;
+      }
+    }
 
-      PyParameter parameter = null;
-      final PsiElement resolvedReference = reference != null? reference.resolve() : null;
-      if (problemElement instanceof PyParameter)
-        parameter = (PyParameter)problemElement;
+    PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
+    if (parameter != null && name != null) {
+      final PyFunction function =
+        elementGenerator.createFromText(LanguageLevel.forElement(problemElement), PyFunction.class,
+                                        "def foo(" + name + ": object):\n\tpass");
+      final PyNamedParameter namedParameter = function.getParameterList().findParameterByName(name);
+      assert namedParameter != null;
+      parameter = (PyParameter)parameter.replace(namedParameter);
+      parameter = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(parameter);
+      editor.getCaretModel().moveToOffset(parameter.getTextOffset());
+
+      final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(parameter);
+      builder.replaceRange(TextRange.create(parameter.getTextLength()-PyNames.OBJECT.length(), parameter.getTextLength()), PyNames.OBJECT);
+      Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
+      TemplateManager.getInstance(project).startTemplate(editor, template);
+    }
+    else {    //return type
+      Callable callable = null;
+      if (resolvedReference instanceof PyTargetExpression) {
+        final PyExpression assignedValue = ((PyTargetExpression)resolvedReference).findAssignedValue();
+        if (assignedValue instanceof PyCallExpression) {
+          callable = ((PyCallExpression)assignedValue).resolveCalleeFunction(PyResolveContext.defaultContext());
+        }
+      }
       else {
-        if (resolvedReference instanceof PyParameter) {
-          parameter = (PyParameter)resolvedReference;
-        }
-      }
-      if (parameter != null && name != null) {
-        final PyFunction function =
-          elementGenerator.createFromText(LanguageLevel.forElement(problemElement), PyFunction.class,
-                                          "def foo(" + name + ": object):\n\tpass");
-        final PyNamedParameter namedParameter = function.getParameterList().findParameterByName(name);
-        assert namedParameter != null;
-        parameter = (PyParameter)parameter.replace(namedParameter);
-        parameter = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(parameter);
-        editor.getCaretModel().moveToOffset(parameter.getTextOffset());
+        PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1);
+        if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER))
+          elementAt = file.findElementAt(editor.getCaretModel().getOffset());
 
-        final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(parameter);
-        builder.replaceRange(TextRange.create(parameter.getTextLength()-PyNames.OBJECT.length(), parameter.getTextLength()), PyNames.OBJECT);
+        callable = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
+      }
+      if (callable instanceof PyFunction && ((PyFunction)callable).getAnnotation() == null) {
+        final String functionSignature = "def " + callable.getName() + callable.getParameterList().getText();
+        String functionText = functionSignature +
+                            " -> object:";
+        for (PyStatement st : ((PyFunction)callable).getStatementList().getStatements()) {
+          functionText = functionText + "\n\t" + st.getText();
+        }
+        final PyFunction function = elementGenerator.createFromText(LanguageLevel.forElement(callable), PyFunction.class,
+                                                                    functionText);
+        callable = (PyFunction)callable.replace(function);
+        callable = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(callable);
+
+        final PyExpression value = ((PyFunction)callable).getAnnotation().getValue();
+        final int offset = value.getTextOffset();
+
+        final TemplateBuilder builder = TemplateBuilderFactory.getInstance().
+          createTemplateBuilder(value);
+        builder.replaceRange(TextRange.create(0, PyNames.OBJECT.length()), PyNames.OBJECT);
         Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
-        TemplateManager.getInstance(project).startTemplate(editor, template);
-      }
-      else {    //return type
-        if (resolvedReference instanceof PyTargetExpression) {
-          final PyExpression assignedValue = ((PyTargetExpression)resolvedReference).findAssignedValue();
-          if (assignedValue instanceof PyCallExpression) {
-            Callable callable = ((PyCallExpression)assignedValue).resolveCalleeFunction(PyResolveContext.defaultContext());
-            if (callable instanceof PyFunction && ((PyFunction)callable).getAnnotation() == null) {
-              final String functionSignature = "def " + callable.getName() + callable.getParameterList().getText();
-              String functionText = functionSignature +
-                                  " -> object:";
-              for (PyStatement st : ((PyFunction)callable).getStatementList().getStatements()) {
-                functionText = functionText + "\n\t" + st.getText();
-              }
-              final PyFunction function = elementGenerator.createFromText(LanguageLevel.forElement(problemElement), PyFunction.class,
-                                                                          functionText);
-              callable = (PyFunction)callable.replace(function);
-              callable = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(callable);
-
-              final PyExpression value = ((PyFunction)callable).getAnnotation().getValue();
-              final int offset = value.getTextOffset();
-
-              final TemplateBuilder builder = TemplateBuilderFactory.getInstance().
-                createTemplateBuilder(value);
-              builder.replaceRange(TextRange.create(0, PyNames.OBJECT.length()), PyNames.OBJECT);
-              Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
-              OpenFileDescriptor descriptor = new OpenFileDescriptor(
-                project,
-                value.getContainingFile().getVirtualFile(),
-                offset
-              );
-              Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
-              if (targetEditor != null) {
-                targetEditor.getCaretModel().moveToOffset(offset);
-                TemplateManager.getInstance(project).startTemplate(targetEditor, template);
-              }
-            }
-          }
+        OpenFileDescriptor descriptor = new OpenFileDescriptor(
+          project,
+          value.getContainingFile().getVirtualFile(),
+          offset
+        );
+        Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
+        if (targetEditor != null) {
+          targetEditor.getCaretModel().moveToOffset(offset);
+          TemplateManager.getInstance(project).startTemplate(targetEditor, template);
         }
-
       }
     }
   }

python/testData/intentions/afterTypeInDocstring1.py

@@ -0,0 +1,5 @@
+def foo(a, b):
+  """
+  :rtype : object
+  """
+  b = 1

python/testData/intentions/beforeTypeInDocstring1.py

@@ -0,0 +1,2 @@
+def f<caret>oo(a, b):
+  b = 1

python/testSrc/com/jetbrains/python/PyIntentionTest.java

@@ -237,6 +237,10 @@ public class PyIntentionTest extends PyTestCase {
     doTest(PyBundle.message("INTN.specify.type"));
   }
 
+  public void testTypeInDocstring1() {
+    doTest(PyBundle.message("INTN.specify.return.type"));
+  }
+
   public void testTypeInPy3Annotation() {      //PY-7045
     doTest(PyBundle.message("INTN.specify.type.in.annotation"), LanguageLevel.PYTHON32);
   }