PY-7124

python/src/com/jetbrains/python/actions/DocstringQuickFix.java

@@ -77,13 +77,14 @@ public class DocstringQuickFix implements LocalQuickFix {
     PyStringLiteralExpression docStringExpression = docStringOwner.getDocStringExpression();
     if (docStringExpression == null && myMissing == null && myUnexpected == null) {
       if (docStringOwner instanceof PyFunction) {
-        PythonDocumentationProvider.inserDocStub((PyFunction)docStringOwner, project, getEditor(project, docStringOwner.getContainingFile()));
+        PythonDocumentationProvider.insertDocStub((PyFunction)docStringOwner, project,
+                                                  getEditor(project, docStringOwner.getContainingFile()));
       }
       if (docStringOwner instanceof PyClass) {
         PyFunction init = ((PyClass)docStringOwner).findInitOrNew(false);
         if (init == null) return;
-        PythonDocumentationProvider.inserDocStub(init, ((PyClass)docStringOwner).getStatementList(),
-                                                 project, getEditor(project, docStringOwner.getContainingFile()));
+        PythonDocumentationProvider.insertDocStub(init, ((PyClass)docStringOwner).getStatementList(),
+                                                  project, getEditor(project, docStringOwner.getContainingFile()));
       }
       return;
     }

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

@@ -52,9 +52,6 @@ public class PyDocStubIntention extends BaseIntentionAction {
   public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
     if (!CodeInsightUtilBase.preparePsiElementForWrite(file)) return;
     PyFunction function = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), PyFunction.class);
-    PythonDocumentationProvider.inserDocStub(function, project, editor);
-
-
+    PythonDocumentationProvider.insertDocStub(function, project, editor);
   }
-
 }

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

@@ -492,7 +492,7 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i
     return result;
   }
 
-  public static void inserDocStub(PyFunction function, PyStatementList insertPlace, Project project, Editor editor) {
+  public static void insertDocStub(PyFunction function, PyStatementList insertPlace, Project project, Editor editor) {
     PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
     PsiWhiteSpace whitespace = PsiTreeUtil.getPrevSiblingOfType(insertPlace, PsiWhiteSpace.class);
     String ws = "\n";
@@ -504,10 +504,16 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i
     String docContent = ws + generateDocumentationContentStub(function, ws, true);
     PyExpressionStatement string = elementGenerator.createDocstring("\"\"\"" + docContent + "\"\"\"");
     if (insertPlace.getStatements().length != 0) {
-      PyFunction func = elementGenerator.createFromText(LanguageLevel.forElement(function),
-                                                        PyFunction.class, "def " + function.getName() + function.getParameterList().getText()
-                                                        +":\n\t"+ string.getText() + "\n\t" + insertPlace.getText());
-      function.replace(func);
+      if (!insertPlace.getText().contains("\n")) {
+        PyFunction func = elementGenerator.createFromText(LanguageLevel.forElement(function),
+                                                          PyFunction.class,
+                                                          "def " + function.getName() + function.getParameterList().getText()
+                                                          + ":\n\t" + string.getText() + "\n\t" + insertPlace.getText());
+        function.replace(func);
+      }
+      else {
+        insertPlace.addBefore(string, insertPlace.getStatements()[0]);
+      }
     }
     PyStringLiteralExpression docstring = function.getDocStringExpression();
     if (editor != null && docstring != null) {
@@ -517,8 +523,8 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i
     }
   }
 
-  public static void inserDocStub(PyFunction function, Project project, Editor editor) {
-    inserDocStub(function, function.getStatementList(), project, editor);
+  public static void insertDocStub(PyFunction function, Project project, Editor editor) {
+    insertDocStub(function, function.getStatementList(), project, editor);
   }
 
   public String generateDocumentationContentStub(PyFunction element, boolean checkReturn) {

python/testData/intentions/afterDocStub.py

@@ -5,5 +5,7 @@ def foo(a, b):
     :param b:
     :return:
     """
+    print a
+    print b
     if True:
         return

python/testData/intentions/afterOneLineDocStub.py

@@ -0,0 +1,7 @@
+def foo(a, b):
+    """
+  
+    :param a:
+    :param b:
+    """
+    pass

python/testData/intentions/beforeDocStub.py

@@ -1,3 +1,5 @@
 def fo<caret>o(a, b):
+    print a
+    print b
     if True:
         return

python/testData/intentions/beforeOneLineDocStub.py

@@ -0,0 +1 @@
+def fo<caret>o(a, b): pass

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

@@ -250,6 +250,14 @@ public class PyIntentionTest extends PyTestCase {
   }
 
   public void testDocStub() {
+    doDocStubTest();
+  }
+
+  public void testOneLineDocStub() {
+    doDocStubTest();
+  }
+
+  private void doDocStubTest() {
     CodeInsightSettings codeInsightSettings = CodeInsightSettings.getInstance();
     codeInsightSettings.JAVADOC_STUB_ON_ENTER = true;
     PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(myFixture.getProject());