fixed PY-6133 Move statement: strange comment moving down to def or class statements

python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/StatementMover.java

@@ -217,17 +217,24 @@ public class StatementMover extends LineMover {
       parent2 = PsiTreeUtil.getParentOfType(element2, PyFunction.class);
       PyElement parent1 = PsiTreeUtil.getParentOfType(myStatementToMove, PyFunction.class);
       if (parent2 != null && parent2 != parent1) {
-        TextRange textRange2 = parent2.getTextRange();
-        int endOffset = myStatementToMove instanceof PsiComment? doc.getLineNumber(textRange2.getEndOffset())
-                                                               : doc.getLineNumber(textRange2.getEndOffset())+1;
+        final TextRange textRange2 = parent2.getTextRange();
+        if (myStatementToMove instanceof PsiComment) {
+          myStatementToIncreaseIndent = myStatementToMove;
+          return;
+        }
+        final int endOffset = doc.getLineNumber(textRange2.getEndOffset())+1;
         info.toMove2 = new LineRange(doc.getLineNumber(textRange2.getStartOffset()), endOffset);
       }
       parent2 = PsiTreeUtil.getParentOfType(element2, PyClass.class);
       parent1 = PsiTreeUtil.getParentOfType(myStatementToMove, PyClass.class);
       if (parent2 != null && parent2 != parent1) {
-        TextRange textRange2 = parent2.getTextRange();
-        info.toMove2 = new LineRange(doc.getLineNumber(textRange2.getStartOffset()),
-                                              doc.getLineNumber(textRange2.getEndOffset())+1);
+        final TextRange textRange2 = parent2.getTextRange();
+        if (myStatementToMove instanceof PsiComment) {
+          myStatementToIncreaseIndent = myStatementToMove;
+          return;
+        }
+        final int endOffset = doc.getLineNumber(textRange2.getEndOffset())+1;
+        info.toMove2 = new LineRange(doc.getLineNumber(textRange2.getStartOffset()), endOffset);
       }
     }
   }

python/testData/mover/commentIntoCompound.py

@@ -0,0 +1,6 @@
+#comment   <caret>
+def f():
+    if True:
+        a = 1
+    else:
+        a = 2

python/testData/mover/commentIntoCompound_afterDown.py

@@ -0,0 +1,6 @@
+def f():
+    #comment
+    if True:
+        a = 1
+    else:
+        a = 2

python/testData/mover/commentIntoCompound_afterUp.py

@@ -0,0 +1,6 @@
+#comment   <caret>
+def f():
+    if True:
+        a = 1
+    else:
+        a = 2

python/testData/mover/commentUp_afterUp.py

@@ -1,6 +1,6 @@
-#comment
 def f():
     if True:
         a = 1
     else:
         a = 2
+        #comment

python/testSrc/com/jetbrains/python/PyStatementMoverTest.java

@@ -161,6 +161,10 @@ public class PyStatementMoverTest extends PyTestCase {
     doTest();
   }
 
+  public void testCommentIntoCompound() { // PY-6133
+    doTest();
+  }
+
   public void testWith() { // PY-5202
     try {
       setLanguageLevel(LanguageLevel.PYTHON27);