fixed PY-7657 Specify type in docstring: AE at com.jetbrains.python.codeInsight.intentions.SpecifyTypeInDocstringIntention.invoke

python/src/com/jetbrains/python/documentation/PyDocstringGenerator.java

@@ -19,8 +19,11 @@ import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.codeStyle.CodeStyleSettings;
+import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.PyNames;
+import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.debugger.PySignatureUtil;
 import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.NotNull;
@@ -188,7 +191,12 @@ public class PyDocstringGenerator {
 
     PsiWhiteSpace whitespace = null;
     if (myDocStringOwner instanceof PyFunction) {
-      whitespace = PsiTreeUtil.getPrevSiblingOfType(((PyFunction)myDocStringOwner).getStatementList(), PsiWhiteSpace.class);
+      final PyStatementList statementList = ((PyFunction)myDocStringOwner).getStatementList();
+      final Document document = PsiDocumentManager.getInstance(myProject).getDocument(getFile());
+      if (document != null && statementList != null && statementList.getStatements().length != 0
+          && document.getLineNumber(statementList.getTextOffset()) != document.getLineNumber(myFunction.getTextOffset())) {
+        whitespace = PsiTreeUtil.getPrevSiblingOfType(statementList, PsiWhiteSpace.class);
+      }
     }
     String ws = "\n";
     if (whitespace != null) {
@@ -197,6 +205,9 @@ public class PyDocstringGenerator {
         ws += whitespace.getText().split("\n")[1];
       }
     }
+    else {
+      ws += StringUtil.repeat(" ", getIndentSize(myFunction));
+    }
     if (replacementText.length() > 0) {
       replacementText.deleteCharAt(replacementText.length() - 1);
     }
@@ -320,12 +331,15 @@ public class PyDocstringGenerator {
       final PyStatementList list = myFunction.getStatementList();
       final Document document = PsiDocumentManager.getInstance(myProject).getDocument(getFile());
 
-      if (list != null && list.getStatements().length != 0) {
-        if (document.getLineNumber(list.getTextOffset()) == document.getLineNumber(myFunction.getTextOffset())) {
+      if (document != null && list != null) {
+        if (document.getLineNumber(list.getTextOffset()) == document.getLineNumber(myFunction.getTextOffset()) ||
+          list.getStatements().length == 0) {
           PyFunction func = elementGenerator.createFromText(LanguageLevel.forElement(myFunction),
-                                                            PyFunction.class,
-                                                            "def " + myFunction.getName() + myFunction.getParameterList().getText()
-                                                            + ":\n\t" + replacementToOffset.getFirst() + "\n\t" + list.getText());
+                                                        PyFunction.class,
+                                                        "def " + myFunction.getName() + myFunction.getParameterList().getText()
+                                                        + ":\n" + StringUtil.repeat(" ",  getIndentSize(myFunction))
+                                                        + replacementToOffset.getFirst() + "\n" +
+                                                        StringUtil.repeat(" ", getIndentSize(myFunction)) + list.getText());
 
           myFunction = (PyFunction)myFunction.replace(func);
         }
@@ -340,6 +354,19 @@ public class PyDocstringGenerator {
     }
   }
 
+  private int getIndentSize(PyFunction function) {
+    CodeStyleSettings.IndentOptions indentOptions = CodeStyleSettingsManager.
+      getInstance(function.getProject()).getCurrentSettings().getIndentOptions(PythonFileType.INSTANCE);
+
+    PyStatementList statementList = PsiTreeUtil.getParentOfType(function, PyStatementList.class);
+    int indent = 1;
+    while (statementList != null) {
+      statementList = PsiTreeUtil.getParentOfType(statementList, PyStatementList.class);
+      indent += 1;
+    }
+    return indent * indentOptions.TAB_SIZE;
+  }
+
   private String getPrefix() {
     PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(myProject);
     String prefix = ":";

python/testData/intentions/afterTypeInDocstring4.py

@@ -0,0 +1,6 @@
+class ProjectElement(object):
+    def __init__(self, project_name='', info_source='',
+                 project_id=None, has_revisions=True):
+        """
+        :type project_name: object
+        """

python/testData/intentions/beforeTypeInDocstring4.py

@@ -0,0 +1,3 @@
+class ProjectElement(object):
+    def __init__(self, pro<caret>ject_name='', info_source='',
+                 project_id=None, has_revisions=True):

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

@@ -234,10 +234,18 @@ public class PyIntentionTest extends PyTestCase {
   }
 
   public void testTypeInDocstring() {
-    doTest(PyBundle.message("INTN.specify.type"));
+    doDocReferenceTest();
   }
 
   public void testTypeInDocstring3() {
+    doDocReferenceTest();
+  }
+
+  public void testTypeInDocstring4() {
+    doDocReferenceTest();
+  }
+
+  private void doDocReferenceTest() {
     doTest(PyBundle.message("INTN.specify.type"));
   }