PY-2903 Take into account preceding comments when picking place for extracted method

Interestingly, test existed already but was disabled a long ago.

python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -18,6 +18,7 @@ package com.jetbrains.python.psi.impl;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.util.Condition;
+import com.intellij.openapi.util.Couple;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.*;
 import com.intellij.psi.stubs.StubElement;
@@ -246,6 +247,37 @@ public class PyPsiUtils {
     }
   }
 
+  /**
+   * Returns comments preceding given elements as pair of the first and the last such comments. Comments should not be
+   * separated by any empty line.
+   * @param element element comments should be adjacent to
+   * @return described range or {@code null} if there are no such comments
+   */
+  @Nullable
+  public static Couple<PsiComment> getPrecedingComments(@NotNull PsiElement element) {
+    PsiComment firstComment = null, lastComment = null;
+    overComments:
+    while (true) {
+      int newLinesCount = 0;
+      for (element = element.getPrevSibling(); element instanceof PsiWhiteSpace; element = element.getPrevSibling()) {
+        newLinesCount += StringUtil.getLineBreakCount(element.getText());
+        if (newLinesCount > 1) {
+          break overComments;
+        }
+      }
+      if (element instanceof PsiComment) {
+        if (lastComment == null) {
+          lastComment = (PsiComment)element;
+        }
+        firstComment = (PsiComment)element;
+      }
+      else {
+        break;
+      }
+    }
+    return lastComment == null ? null : Couple.of(firstComment, lastComment);
+  }
+
   @NotNull
   static <T, U extends PsiElement> List<T> collectStubChildren(U e,
                                                                final StubElement<U> stub, final IElementType elementType,

python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java

@@ -24,6 +24,7 @@ import com.intellij.openapi.command.CommandProcessor;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.Messages;
+import com.intellij.openapi.util.Couple;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.StringUtil;
@@ -466,15 +467,17 @@ public class PyExtractMethodUtil {
     }
     final PsiNamedElement parent = PsiTreeUtil.getParentOfType(anchor, PyFile.class, PyClass.class, PyFunction.class);
 
-    final PsiElement result;
+    final PsiElement insertionAnchor;
     if (parent instanceof PyFile || parent instanceof PyClass) {
       final PsiElement target = parent instanceof PyClass ? ((PyClass)parent).getStatementList() : parent;
-      final PsiElement anchorStatement = PyPsiUtils.getParentRightBefore(anchor, target);
-      result = target.addBefore(generatedMethod, anchorStatement);
+      insertionAnchor = PyPsiUtils.getParentRightBefore(anchor, target);
     }
     else {
-      result = parent.getParent().addBefore(generatedMethod, parent);
+      insertionAnchor = parent;
     }
+    assert insertionAnchor != null;
+    final Couple<PsiComment> comments = PyPsiUtils.getPrecedingComments(insertionAnchor);
+    final PsiElement result = insertionAnchor.getParent().addBefore(generatedMethod, comments != null ? comments.getFirst() : insertionAnchor);
     // to ensure correct reformatting, mark the entire method as generated
     result.accept(new PsiRecursiveElementVisitor() {
       @Override

python/testData/refactoring/extractmethod/SeveralCommentsAbove.after.py

@@ -0,0 +1,9 @@
+# Unrelated comment
+def bar():
+    print("Hello")
+
+
+# Comment to method
+# continuation
+def foo():
+    bar()

python/testData/refactoring/extractmethod/SeveralCommentsAbove.before.py

@@ -0,0 +1,6 @@
+# Unrelated comment
+
+# Comment to method
+# continuation
+def foo():
+    <selection>print("Hello")</selection>

python/testSrc/com/jetbrains/python/refactoring/PyExtractMethodTest.java

@@ -122,7 +122,12 @@ public class PyExtractMethodTest extends LightMarkedTestCase {
   }
 
   // PY-2903
-  public void _testComment() {
+  public void testComment() {
+    doTest("bar");
+  }
+
+  // PY-2903
+  public void testSeveralCommentsAbove() {
     doTest("bar");
   }