From 7fd1c618c0102554a0cd908f581ee41e84375110 Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Mon, 8 Sep 2014 19:16:48 +0400
Subject: [PATCH] fixed PY-11198 Move Statement Up/Down near dictionary
 expression unexpectedly merges statements

---
 .../moveUpDown/PyStatementMover.java          | 21 ++++++++++++++-----
 .../src/com/jetbrains/python/psi/PyUtil.java  |  4 ++--
 python/testData/mover/sameLevelAsDict.py      |  9 ++++++++
 .../mover/sameLevelAsDict_afterDown.py        |  9 ++++++++
 .../testData/mover/sameLevelAsDict_afterUp.py |  9 ++++++++
 .../python/PyStatementMoverTest.java          |  6 +++++-
 6 files changed, 50 insertions(+), 8 deletions(-)
 create mode 100644 python/testData/mover/sameLevelAsDict.py
 create mode 100644 python/testData/mover/sameLevelAsDict_afterDown.py
 create mode 100644 python/testData/mover/sameLevelAsDict_afterUp.py

diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/PyStatementMover.java b/python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/PyStatementMover.java
index 06b882b2066c..89b6d9bc4a5d 100644
--- a/python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/PyStatementMover.java
+++ b/python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/PyStatementMover.java
@@ -108,7 +108,6 @@ public class PyStatementMover extends LineMover {
     if (moveOutsideFile(document, lineNumber)) return null;
     int lineEndOffset = document.getLineEndOffset(lineNumber);
     final int startOffset = document.getLineStartOffset(lineNumber);
-    lineEndOffset = startOffset != lineEndOffset ? lineEndOffset - 1 : lineEndOffset;
 
     final PyStatementList statementList = getStatementList(elementToMove);
 
@@ -278,11 +277,23 @@ public class PyStatementMover extends LineMover {
 
   private static PsiElement getDestinationElement(@NotNull final PsiElement elementToMove, @NotNull final Document document,
                                                   int lineEndOffset, boolean down) {
-    PsiElement destination = elementToMove.getContainingFile().findElementAt(lineEndOffset);
-    if (destination == null) return null;
-    if (destination instanceof PsiComment) return destination;
+    PsiElement destination = PyUtil.findPrevAtOffset(elementToMove.getContainingFile(), lineEndOffset, PsiWhiteSpace.class);
     PsiElement sibling = down ? PsiTreeUtil.getNextSiblingOfType(elementToMove, PyStatement.class) :
-                  PsiTreeUtil.getPrevSiblingOfType(elementToMove, PyStatement.class);
+                         PsiTreeUtil.getPrevSiblingOfType(elementToMove, PyStatement.class);
+    if (destination == null) {
+      if (elementToMove instanceof PyClass) {
+        destination = sibling;
+      }
+      else if (elementToMove instanceof PyFunction) {
+        if (!(sibling instanceof PyClass))
+          destination = sibling;
+        else destination = null;
+      }
+      else {
+        return null;
+      }
+    }
+    if (destination instanceof PsiComment) return destination;
     if (elementToMove instanceof PyClass) {
       destination = sibling;
     }
diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java
index cd4e582bb6f2..dc5cbf4917a5 100644
--- a/python/src/com/jetbrains/python/psi/PyUtil.java
+++ b/python/src/com/jetbrains/python/psi/PyUtil.java
@@ -1394,8 +1394,8 @@ public class PyUtil {
 
   @Nullable
   public static PsiElement findPrevAtOffset(PsiFile psiFile, int caretOffset, Class... toSkip) {
-    PsiElement element = psiFile.findElementAt(caretOffset);
-    if (element == null || caretOffset < 0) {
+    PsiElement element;
+    if (caretOffset < 0) {
       return null;
     }
     int lineStartOffset = 0;
diff --git a/python/testData/mover/sameLevelAsDict.py b/python/testData/mover/sameLevelAsDict.py
new file mode 100644
index 000000000000..fd5b54ec4e82
--- /dev/null
+++ b/python/testData/mover/sameLevelAsDict.py
@@ -0,0 +1,9 @@
+a = {
+    'c': 99999
+}
+print "Hello, there."
+a = {
+    'b': 1,
+    'c': 2
+}
+print <caret>a['c']
\ No newline at end of file
diff --git a/python/testData/mover/sameLevelAsDict_afterDown.py b/python/testData/mover/sameLevelAsDict_afterDown.py
new file mode 100644
index 000000000000..602da928bab7
--- /dev/null
+++ b/python/testData/mover/sameLevelAsDict_afterDown.py
@@ -0,0 +1,9 @@
+a = {
+    'c': 99999
+}
+print "Hello, there."
+a = {
+    'b': 1,
+    'c': 2
+}
+print a['c']
\ No newline at end of file
diff --git a/python/testData/mover/sameLevelAsDict_afterUp.py b/python/testData/mover/sameLevelAsDict_afterUp.py
new file mode 100644
index 000000000000..c4817e7a6474
--- /dev/null
+++ b/python/testData/mover/sameLevelAsDict_afterUp.py
@@ -0,0 +1,9 @@
+a = {
+    'c': 99999
+}
+print "Hello, there."
+print a['c']
+a = {
+    'b': 1,
+    'c': 2
+}
diff --git a/python/testSrc/com/jetbrains/python/PyStatementMoverTest.java b/python/testSrc/com/jetbrains/python/PyStatementMoverTest.java
index 86a4728f6ce1..220c08ef2d24 100644
--- a/python/testSrc/com/jetbrains/python/PyStatementMoverTest.java
+++ b/python/testSrc/com/jetbrains/python/PyStatementMoverTest.java
@@ -260,7 +260,11 @@ public class PyStatementMoverTest extends PyTestCase {
     doTest();
   }
 
-  public void testOutsideFromDict() { //PY-11595
+  public void testOutsideFromDict() {
+    doTest();
+  }
+
+  public void testSameLevelAsDict() {
     doTest();
   }