PY-3032 Surround action can't be invoked when first line of selection is a comment

python/src/com/jetbrains/python/refactoring/PyRefactoringUtil.java

@@ -25,7 +25,6 @@ import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtilCore;
 import com.intellij.usageView.UsageInfo;
-import com.intellij.util.Processor;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.HashSet;
 import com.jetbrains.python.PyNames;
@@ -204,16 +203,35 @@ public class PyRefactoringUtil {
 
   @NotNull
   public static PsiElement[] findStatementsInRange(@NotNull final PsiFile file, int startOffset, int endOffset) {
+    ArrayList<PsiElement> array = new ArrayList<>();
+
     PsiElement element1 = file.findElementAt(startOffset);
     PsiElement element2 = file.findElementAt(endOffset - 1);
+    PsiElement endComment = null;
+
+    boolean startsWithWhitespace = false;
+    boolean endsWithWhitespace = false;
     if (element1 instanceof PsiWhiteSpace) {
       startOffset = element1.getTextRange().getEndOffset();
       element1 = file.findElementAt(startOffset);
+      startsWithWhitespace = true;
     }
     if (element2 instanceof PsiWhiteSpace) {
-      endOffset = element2.getTextRange().getStartOffset();
-      element2 = file.findElementAt(endOffset - 1);
+      element2 = PsiTreeUtil.skipSiblingsBackward(element2, PsiWhiteSpace.class);
+      endsWithWhitespace = true;
     }
+    while (element2 instanceof PsiComment) {
+      endComment = element2;
+      element2 = PsiTreeUtil.skipSiblingsBackward(element2, PsiWhiteSpace.class, PsiComment.class);
+      endsWithWhitespace = true;
+    }
+
+    while (element1 instanceof PsiComment) {
+      array.add(element1);
+      element1 = PsiTreeUtil.skipSiblingsForward(element1, PsiWhiteSpace.class);
+      startsWithWhitespace = true;
+    }
+
     if (element1 == null || element2 == null) {
       return PsiElement.EMPTY_ARRAY;
     }
@@ -242,7 +260,7 @@ public class PyRefactoringUtil {
         element1 = element1.getParent();
       }
     }
-    if (startOffset != element1.getTextRange().getStartOffset()) {
+    if (startOffset != element1.getTextRange().getStartOffset() && !startsWithWhitespace) {
       return PsiElement.EMPTY_ARRAY;
     }
 
@@ -251,7 +269,7 @@ public class PyRefactoringUtil {
         element2 = element2.getParent();
       }
     }
-    if (endOffset != element2.getTextRange().getEndOffset()) {
+    if (endOffset != element2.getTextRange().getEndOffset() && !endsWithWhitespace) {
       return PsiElement.EMPTY_ARRAY;
     }
 
@@ -263,7 +281,7 @@ public class PyRefactoringUtil {
     }
 
     PsiElement[] children = parent.getChildren();
-    ArrayList<PsiElement> array = new ArrayList<>();
+
     boolean flag = false;
     for (PsiElement child : children) {
       if (child.equals(element1)) {
@@ -277,6 +295,11 @@ public class PyRefactoringUtil {
       }
     }
 
+    while (endComment instanceof PsiComment) {
+      array.add(endComment);
+      endComment = PsiTreeUtil.skipSiblingsForward(endComment, PsiWhiteSpace.class);
+    }
+
     for (PsiElement element : array) {
       if (!(element instanceof PyStatement || element instanceof PsiWhiteSpace || element instanceof PsiComment)) {
         return PsiElement.EMPTY_ARRAY;

python/testData/surround/SurroundCommentAtEnd.py

@@ -0,0 +1,2 @@
+<selection>a = 1
+#comment line</selection>

python/testData/surround/SurroundCommentAtEnd_after.py

@@ -0,0 +1,3 @@
+if True:
+    a = 1
+    #comment line

python/testData/surround/SurroundCommentAtStart.py

@@ -0,0 +1,2 @@
+<selection>#comment line
+a = 1</selection>

python/testData/surround/SurroundCommentAtStart_after.py

@@ -0,0 +1,3 @@
+if True:
+    #comment line
+    a = 1

python/testData/surround/SurroundNewline.py

@@ -0,0 +1,3 @@
+<selection>a = 1
+</selection>
+a = 2

python/testData/surround/SurroundNewline_after.py

@@ -0,0 +1,4 @@
+if True:
+    a = 1
+
+a = 2

python/testSrc/com/jetbrains/python/PySurroundWithTest.java

@@ -83,6 +83,18 @@ public class PySurroundWithTest extends PyTestCase {
     doTest(surrounder);
   }
 
+  public void testSurroundCommentAtStart() throws Exception {
+    doTest(new PyWithIfSurrounder());
+  }
+
+  public void testSurroundCommentAtEnd() throws Exception {
+    doTest(new PyWithIfSurrounder());
+  }
+
+  public void testSurroundNewline() throws Exception {
+    doTest(new PyWithIfSurrounder());
+  }
+
   private void doTest(final Surrounder surrounder) throws Exception {
     String baseName = "/surround/" + getTestName(false);
     myFixture.configureByFile(baseName + ".py");