PY-17183 Heuristically determine when incomplete docstring captures another declaration

To do so I search for the line inside a docstring that starts with
either "def " or "class " and has indentation less than that of opening
docstring quotes.

python/src/com/jetbrains/python/editor/PythonEnterHandler.java

@@ -29,6 +29,7 @@ import com.intellij.openapi.editor.actionSystem.EditorActionHandler;
 import com.intellij.openapi.editor.actions.SplitLineAction;
 import com.intellij.openapi.util.Ref;
 import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.LineTokenizer;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.source.tree.TreeUtil;
 import com.intellij.psi.impl.source.tree.injected.InjectedLanguageUtil;
@@ -106,7 +107,7 @@ public class PythonEnterHandler extends EnterHandlerDelegateAdapter {
         comment = file.findElementAt(offset - 1);
       }
       int expectedStringStart = editor.getCaretModel().getOffset() - 3; // """ or '''
-      if (comment != null && atDocCommentStart(comment, expectedStringStart)) {
+      if (comment != null && atDocCommentStart(comment, expectedStringStart, doc)) {
         insertDocStringStub(editor, comment);
         return Result.Continue;
       }
@@ -379,14 +380,22 @@ public class PythonEnterHandler extends EnterHandlerDelegateAdapter {
     }
   }
 
-  public static boolean atDocCommentStart(@NotNull PsiElement element, int offset) {
+  public static boolean atDocCommentStart(@NotNull PsiElement element, int firstQuoteOffset, @NotNull Document document) {
+    if (firstQuoteOffset < 0 || firstQuoteOffset > document.getTextLength() - 3) {
+      return false;
+    }
+    final String quotes = document.getText(TextRange.from(firstQuoteOffset, 3));
+    if (!quotes.equals("\"\"\"") && !quotes.equals("'''")) {
+      return false;
+    }
     final PyStringLiteralExpression pyString = DocStringUtil.getParentDefinitionDocString(element);
     if (pyString != null) {
-      String text = element.getText();
-      final int prefixLength = PyStringLiteralExpressionImpl.getPrefixLength(text);
-      text = text.substring(prefixLength);
-      if (pyString.getText().endsWith(text) && (text.startsWith("\"\"\"") || text.startsWith("'''"))) {
-        if (offset == pyString.getTextOffset() + prefixLength) {
+      String nodeText = element.getText();
+      final int prefixLength = PyStringLiteralExpressionImpl.getPrefixLength(nodeText);
+      nodeText = nodeText.substring(prefixLength);
+      final String literalText = pyString.getText();
+      if (literalText.endsWith(nodeText) && nodeText.startsWith(quotes)) {
+        if (firstQuoteOffset == pyString.getTextOffset() + prefixLength) {
           PsiErrorElement error = PsiTreeUtil.getNextSiblingOfType(pyString, PsiErrorElement.class);
           if (error != null) {
             return true;
@@ -396,9 +405,20 @@ public class PythonEnterHandler extends EnterHandlerDelegateAdapter {
             return true;
           }
 
-          if (text.length() < 6 || (!text.endsWith("\"\"\"") && !text.endsWith("'''"))) {
+          if (nodeText.length() < 6 || !nodeText.endsWith(quotes)) {
             return true;
           }
+          // Sometimes if incomplete docstring is followed by another declaration with a docstring, it might be treated
+          // as complete docstring, because we can't understand that closing quotes actually belong to another docstring.
+          final String docstringIndent = PyIndentUtil.getLineIndent(document, document.getLineNumber(firstQuoteOffset));
+          for (String line : LineTokenizer.tokenizeIntoList(nodeText, false)) {
+            final String lineIndent = (String)PyIndentUtil.getLineIndent(line);
+            final String lineContent = line.substring(lineIndent.length());
+            if ((lineContent.startsWith("def ") || lineContent.startsWith("class ")) &&
+                docstringIndent.length() > lineIndent.length() && docstringIndent.startsWith(lineIndent)) {
+              return true;
+            }
+          }
         }
       }
     }

python/src/com/jetbrains/python/editor/PythonSpaceHandler.java

@@ -44,10 +44,10 @@ public class PythonSpaceHandler extends TypedHandlerDelegate {
       }
       if (element == null) return Result.CONTINUE;
       int expectedStringStart = offset - 4;        // """ or ''' plus space char
-      if (PythonEnterHandler.atDocCommentStart(element, expectedStringStart)) {
+      final Document document = editor.getDocument();
+      if (PythonEnterHandler.atDocCommentStart(element, expectedStringStart, document)) {
         final PyDocStringOwner docOwner = PsiTreeUtil.getParentOfType(element, PyDocStringOwner.class);
         if (docOwner != null) {
-          final Document document = editor.getDocument();
           final String quotes = document.getText(TextRange.from(expectedStringStart, 3));
           final String docString = PyDocstringGenerator.forDocStringOwner(docOwner)
             .forceNewMode()

python/src/com/jetbrains/python/psi/PyIndentUtil.java

@@ -16,7 +16,9 @@
 package com.jetbrains.python.psi;
 
 import com.google.common.collect.Iterables;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.LineTokenizer;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
@@ -195,4 +197,11 @@ public class PyIndentUtil {
     }
     return StringUtil.notNullize(minIndent);
   }
+
+  @NotNull
+  public static String getLineIndent(@NotNull Document document, int lineNumber) {
+    final TextRange lineRange = TextRange.create(document.getLineStartOffset(lineNumber), document.getLineEndOffset(lineNumber));
+    final String line = document.getText(lineRange);
+    return (String)getLineIndent(line);
+  }
 }

python/testData/editing/enterDocstringStubWhenClassDocstringBelow.after.py

@@ -0,0 +1,12 @@
+def f():
+    """
+    
+    Returns:
+
+    """
+    
+    
+class Class:
+    """
+    bar
+    """

python/testData/editing/enterDocstringStubWhenClassDocstringBelow.py

@@ -0,0 +1,8 @@
+def f():
+    """<caret>
+    
+    
+class Class:
+    """
+    bar
+    """

python/testData/editing/enterDocstringStubWhenFunctionDocstringBelow.after.py

@@ -0,0 +1,12 @@
+def f():
+    """
+    
+    Returns:
+
+    """
+    
+    
+def g():
+    """
+    bar
+    """

python/testData/editing/enterDocstringStubWhenFunctionDocstringBelow.py

@@ -0,0 +1,8 @@
+def f():
+    """<caret>
+    
+    
+def g():
+    """
+    bar
+    """

python/testData/editing/enterNoDocstringStubWhenCodeExampleInDocstring.after.py

@@ -0,0 +1,15 @@
+def f():
+    """
+    
+    Monospaced ``func`` and func
+
+    Example:
+
+        ::
+        
+            def func():
+                pass
+
+            class Class():
+                pass
+    """

python/testData/editing/enterNoDocstringStubWhenCodeExampleInDocstring.py

@@ -0,0 +1,14 @@
+def f():
+    """<caret>
+    Monospaced ``func`` and func
+
+    Example:
+
+        ::
+        
+            def func():
+                pass
+
+            class Class():
+                pass
+    """

python/testSrc/com/jetbrains/python/PyEditingTest.java

@@ -257,6 +257,21 @@ public class PyEditingTest extends PyTestCase {
     doDocStringTypingTest("\nparam", DocStringFormat.GOOGLE);
   }
 
+  // PY-17183
+  public void testEnterDocstringStubWhenFunctionDocstringBelow() {
+    doDocStringTypingTest("\n", DocStringFormat.GOOGLE);
+  }
+  
+  // PY-17183
+  public void testEnterDocstringStubWhenClassDocstringBelow() {
+    doDocStringTypingTest("\n", DocStringFormat.GOOGLE);
+  }
+
+  // PY-17183
+  public void testEnterNoDocstringStubWhenCodeExampleInDocstring() {
+    doDocStringTypingTest("\n", DocStringFormat.GOOGLE);
+  }
+
   public void testEnterInString() {  // PY-1738
     doTestEnter("a = \"some <caret>string\"", "a = \"some \" \\\n" +
                                               "    \"string\"");