From 03762a6f9d70d96fdd3ea76d85d4f591ea8c80d7 Mon Sep 17 00:00:00 2001
From: "alexey.ivanov" <alexey.ivanov@jetbrains.com>
Date: Tue, 11 May 2010 14:49:17 +0400
Subject: [PATCH] more improvements for python statement mover

---
 .../moveUpDown/StatementMover.java            | 28 +++++++++++++------
 python/testData/mover/indent.py               |  3 +-
 python/testData/mover/indent_afterDown.py     |  5 ++--
 python/testData/mover/indent_afterUp.py       |  3 +-
 4 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/StatementMover.java b/python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/StatementMover.java
index 0c08152c39c0..528c66990856 100644
--- a/python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/StatementMover.java
+++ b/python/src/com/jetbrains/python/codeInsight/editorActions/moveUpDown/StatementMover.java
@@ -224,7 +224,7 @@ public class StatementMover extends LineMover {
         final PyStatementList statementList = parentStatementPart.getStatementList();
         assert statementList != null;
         if (down) {
-          final PyStatementPart nextStatementPart = PsiTreeUtil.getNextSiblingOfType(statementList.getParent(), PyStatementPart.class);
+          PyStatementPart nextStatementPart = PsiTreeUtil.getNextSiblingOfType(statementList.getParent(), PyStatementPart.class);
           if (nextStatementPart != null) {
             info.toMove2 = new LineRange(range.endLine, range.endLine + 1);
             myStatementListToRemovePass = nextStatementPart.getStatementList();
@@ -241,15 +241,20 @@ public class StatementMover extends LineMover {
               if (nextStatement instanceof PyFunction || nextStatement instanceof PyClass) {
                 return false;
               }
-              if (nextStatement == null) {
-                continue;
+              if (nextStatement != null) {
+                break;
               }
-              break;
             }
             final int startLine = editor.offsetToLogicalPosition(parent.getTextRange().getEndOffset()).line;
             final int endLine = editor.offsetToLogicalPosition(nextStatement.getTextRange().getEndOffset()).line;
             info.toMove2 = new LineRange(startLine + 1, endLine + 1);
-            calculateIndent(editor, statementList, nextStatement);
+            nextStatementPart = PsiTreeUtil.getChildOfType(nextStatement, PyStatementPart.class);
+            if (nextStatementPart != null) {
+              calculateIndent(editor, statementList, nextStatementPart.getStatementList(), down);
+            }
+            else {
+              calculateIndent(editor, statementList, nextStatement, down);
+            }
             myElementsToChangeIndent = statements;
           }
         }
@@ -294,7 +299,7 @@ public class StatementMover extends LineMover {
       if (startLineNumber != endLineNumber) {
         info.toMove2 = new LineRange(startLineNumber, startLineNumber + 1);
         myStatementListToRemovePass = statementPart.getStatementList();
-        calculateIndent(editor, range.firstElement, myStatementListToRemovePass);
+        calculateIndent(editor, range.firstElement, myStatementListToRemovePass, down);
         myElementsToChangeIndent = statements;
       }
     }
@@ -302,10 +307,17 @@ public class StatementMover extends LineMover {
   }
 
 
-  private void calculateIndent(final Editor editor, final PsiElement firstElement, final PsiElement secondElement) {
+  private void calculateIndent(final Editor editor, final PsiElement firstElement, final PsiElement secondElement, final boolean down) {
     final PsiFile file = firstElement.getContainingFile();
     final int firstIndent = getIndent(editor, file, editor.offsetToLogicalPosition(firstElement.getTextRange().getStartOffset()).line);
-    final int secondIndent = getIndent(editor, file, editor.offsetToLogicalPosition(secondElement.getTextRange().getEndOffset()).line);
+    int line;
+    if (down) {
+      line = editor.offsetToLogicalPosition(secondElement.getTextRange().getStartOffset()).line;
+    }
+    else {
+      line = editor.offsetToLogicalPosition(secondElement.getTextRange().getEndOffset()).line;
+    }
+    final int secondIndent = getIndent(editor, file, line);
     myIndentLevel = (secondIndent - firstIndent) >> 2;
   }
 
diff --git a/python/testData/mover/indent.py b/python/testData/mover/indent.py
index f9686dab3689..ccf8206ed37a 100644
--- a/python/testData/mover/indent.py
+++ b/python/testData/mover/indent.py
@@ -6,7 +6,8 @@ class A:
                     if d:
                         self.bar()
                 c<caret> = 3
-        a = 2
+        if c:
+            a = 2
 
     def bar(self):
         pass
\ No newline at end of file
diff --git a/python/testData/mover/indent_afterDown.py b/python/testData/mover/indent_afterDown.py
index 695849124e67..87e9a82d4837 100644
--- a/python/testData/mover/indent_afterDown.py
+++ b/python/testData/mover/indent_afterDown.py
@@ -5,8 +5,9 @@ class A:
                 if c:
                     if d:
                         self.bar()
-        a = 2
-        c = 3
+        if c:
+            a = 2
+            c = 3
 
     def bar(self):
         pass
\ No newline at end of file
diff --git a/python/testData/mover/indent_afterUp.py b/python/testData/mover/indent_afterUp.py
index dae668039366..45ee716ff217 100644
--- a/python/testData/mover/indent_afterUp.py
+++ b/python/testData/mover/indent_afterUp.py
@@ -6,7 +6,8 @@ class A:
                     if d:
                         c = 3
                         self.bar()
-        a = 2
+        if c:
+            a = 2
 
     def bar(self):
         pass
\ No newline at end of file