PY-41375 PY-37765 Fix selection problem in PyStatementMover

Also fix unchecked exception in PY-37765

(cherry picked from commit 3b8f3ca228d8e649a63d277314bf13994b80ffd3)

IJ-MR-6623

GitOrigin-RevId: 22026406e912c057bb7d81cb59d7a4e21da6f0f5

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

@@ -341,6 +341,8 @@ public class PyStatementMover extends LineMover {
         final CaretModel caretModel = editor.getCaretModel();
 
         final int selectionStart = selectionModel.getSelectionStart();
+        final LogicalPosition selectionBeforeMoveStart = editor.offsetToLogicalPosition(selectionStart);
+        final LogicalPosition selectionBeforeMoveEnd = editor.offsetToLogicalPosition(selectionModel.getSelectionEnd());
         boolean isSelectionStartAtCaret = caretModel.getOffset() == selectionStart;
         final SelectionContainer selectionLen = getSelectionLenContainer(editor, ((MyLineRange)toMove));
 
@@ -354,8 +356,12 @@ public class PyStatementMover extends LineMover {
         else {
           offset = moveInOut(((MyLineRange)toMove), editor, info);
         }
+        final LogicalPosition positionOffsetAfterMove = editor.offsetToLogicalPosition(offset);
         restoreCaretAndSelection(file, editor, isSelectionStartAtCaret, hasSelection, selectionLen,
-                                 shift, offset, (MyLineRange)toMove);
+                                 shift, offset, (MyLineRange)toMove,
+                                 selectionBeforeMoveStart,
+                                 selectionBeforeMoveEnd,
+                                 positionOffsetAfterMove);
         info.toMove2 = info.toMove;   //do not move further
       });
     }
@@ -388,7 +394,10 @@ public class PyStatementMover extends LineMover {
 
   private static void restoreCaretAndSelection(@NotNull final PsiFile file, @NotNull final Editor editor, boolean selectionStartAtCaret,
                                                boolean hasSelection, @NotNull final SelectionContainer selectionContainer, int shift,
-                                               int offset, @NotNull final MyLineRange toMove) {
+                                               int offset, @NotNull final MyLineRange toMove,
+                                               LogicalPosition selectionBeforeMoveStart,
+                                               LogicalPosition selectionBeforeMoveEnd,
+                                               LogicalPosition positionOffsetAfterMove) {
     final Document document = editor.getDocument();
     final SelectionModel selectionModel = editor.getSelectionModel();
     final CaretModel caretModel = editor.getCaretModel();
@@ -427,13 +436,31 @@ public class PyStatementMover extends LineMover {
     caretModel.moveToOffset(newCaretOffset);
 
     if (hasSelection) {
-      if (selectionStartAtCaret) {
+      int selectionLinesDiff = selectionBeforeMoveEnd.line - selectionBeforeMoveStart.line;
-        int newSelectionEnd = newCaretOffset + selectionLen;
+      if (selectionLinesDiff > 1) {
-        selectionModel.setSelection(newCaretOffset, newSelectionEnd);
+        int endOffsetColumn = selectionBeforeMoveEnd.column;
+        if (endOffsetColumn > 0) endOffsetColumn += positionOffsetAfterMove.column - selectionBeforeMoveStart.column;
+        int startOffsetColumn = positionOffsetAfterMove.column;
+        if (selectionBeforeMoveStart.column == 0) startOffsetColumn = 0;
+        int startOffset = editor.logicalPositionToOffset(new LogicalPosition(positionOffsetAfterMove.line, startOffsetColumn));
+        int endOffset = editor.logicalPositionToOffset(new LogicalPosition(positionOffsetAfterMove.line + selectionLinesDiff, endOffsetColumn));
+        selectionModel.setSelection(startOffset, endOffset);
+        if (selectionStartAtCaret) {
+          caretModel.moveToOffset(startOffset);
+        }
+        else {
+          caretModel.moveToOffset(endOffset);
+        }
       }
       else {
-        int newSelectionStart = newCaretOffset - selectionLen;
+        if (selectionStartAtCaret) {
-        selectionModel.setSelection(newSelectionStart, newCaretOffset);
+          int newSelectionEnd = newCaretOffset + selectionLen;
+          selectionModel.setSelection(newCaretOffset, newSelectionEnd);
+        }
+        else {
+          int newSelectionStart = newCaretOffset - selectionLen;
+          selectionModel.setSelection(newSelectionStart, newCaretOffset);
+        }
       }
     }
   }

python/testData/mover/multiLineSelectionDifferentLevelsMoveIntoFunction.py

@@ -0,0 +1,11 @@
+import numpy as np
+
+
+def split_by_words(X):
+    X = np.core.chararray.lower(X)
+    return np.core.chararray.split(X)
+<selection>DELIMITERS = "!?:;,.\'-+/\\()"
+
+def parse(string):
+    return "".join((" " if char in DELIMITERS else char) for char in string).split()
+</selection><caret>

python/testData/mover/multiLineSelectionDifferentLevelsMoveIntoFunction_afterDown.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+
+def split_by_words(X):
+    X = np.core.chararray.lower(X)
+    return np.core.chararray.split(X)
+
+<selection>DELIMITERS = "!?:;,.\'-+/\\()"
+
+def parse(string):
+    return "".join((" " if char in DELIMITERS else char) for char in string).split()
+</selection><caret>

python/testData/mover/multiLineSelectionDifferentLevelsMoveIntoFunction_afterUp.py

@@ -0,0 +1,11 @@
+import numpy as np
+
+
+def split_by_words(X):
+    X = np.core.chararray.lower(X)
+    return np.core.chararray.split(X)
+<selection>    DELIMITERS = "!?:;,.\'-+/\\()"
+
+    def parse(string):
+        return "".join((" " if char in DELIMITERS else char) for char in string).split()
+</selection><caret>

python/testData/mover/multiLineSelectionDifferentLevelsMoveOutFromFunction.py

@@ -0,0 +1,4 @@
+def func():
+    <caret><selection>n = 0
+    while n:
+        print("spam")</selection>

python/testData/mover/multiLineSelectionDifferentLevelsMoveOutFromFunction_afterDown.py

@@ -0,0 +1,5 @@
+def func():
+    pass
+<caret><selection>n = 0
+while n:
+    print("spam")</selection>

python/testData/mover/multiLineSelectionDifferentLevelsMoveOutFromFunction_afterUp.py

@@ -0,0 +1,5 @@
+<caret><selection>n = 0
+while n:
+    print("spam")</selection>
+def func():
+    pass

python/testData/mover/multiLineSelectionDifferentLevelsMoveOutFromNestedFunction.py

@@ -0,0 +1,5 @@
+def func():
+    def nested_func():
+        <selection>n = 0
+        while n:
+            print("spam")</selection><caret>

python/testData/mover/multiLineSelectionDifferentLevelsMoveOutFromNestedFunction_afterDown.py

@@ -0,0 +1,6 @@
+def func():
+    def nested_func():
+        pass
+    <selection>n = 0
+    while n:
+        print("spam")</selection><caret>

python/testData/mover/multiLineSelectionDifferentLevelsMoveOutFromNestedFunction_afterUp.py

@@ -0,0 +1,6 @@
+def func():
+    <selection>n = 0
+    while n:
+        print("spam")</selection><caret>
+    def nested_func():
+        pass

python/testData/mover/singleLineSelectionOutFromFunction.py

@@ -0,0 +1,2 @@
+def foo():
+    <selection>a = 42</selection><caret>

python/testData/mover/singleLineSelectionOutFromFunction_afterDown.py

@@ -0,0 +1,3 @@
+def foo():
+    pass
+<selection>a = 42</selection><caret>

python/testData/mover/singleLineSelectionOutFromFunction_afterUp.py

@@ -0,0 +1,3 @@
+<selection>a = 42</selection><caret>
+def foo():
+    pass

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

@@ -248,6 +248,26 @@ public class PyStatementMoverTest extends PyTestCase {
     doTest();
   }
 
+  // PY-41375
+  public void testMultiLineSelectionDifferentLevelsMoveOutFromFunction() {
+    doTest();
+  }
+
+  // PY-37765
+  public void testMultiLineSelectionDifferentLevelsMoveIntoFunction() {
+    doTest();
+  }
+
+  // PY-41375
+  public void testSingleLineSelectionOutFromFunction() {
+    doTest();
+  }
+
+  // PY-41375
+  public void testMultiLineSelectionDifferentLevelsMoveOutFromNestedFunction() {
+    doTest();
+  }
+
   public void testTheSameLevelMultiple() { //PY-10947
     doTest();
   }