PY-296 'Move statement up/down' on method name line should move the whole method

python/src/META-INF/python-plugin-common.xml

@@ -139,6 +139,8 @@
     <codeInsight.gotoSuper language="Python" implementationClass="com.jetbrains.python.codeInsight.PyGotoSuperHandler"/>
     <gotoDeclarationHandler implementation="com.jetbrains.python.codeInsight.PyBreakContinueGotoProvider" order="FIRST"/>
     <lang.smartEnterProcessor language="Python" implementationClass="com.jetbrains.python.codeInsight.editorActions.smartEnter.PySmartEnterProcessor"/>
+    <statementUpDownMover implementation="com.jetbrains.python.codeInsight.editorActions.moveUpDown.StatementMover" id="statement"
+                          order="before line"/>
 
     <filePropertyPusher implementation="com.jetbrains.python.psi.impl.PythonLanguageLevelPusher"/>
 

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

@@ -0,0 +1,342 @@
+package com.jetbrains.python.codeInsight.editorActions.moveUpDown;
+
+import com.intellij.codeInsight.CodeInsightUtilBase;
+import com.intellij.codeInsight.editorActions.moveUpDown.LineMover;
+import com.intellij.codeInsight.editorActions.moveUpDown.LineRange;
+import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.LogicalPosition;
+import com.intellij.openapi.editor.SelectionModel;
+import com.intellij.openapi.util.Pair;
+import com.intellij.psi.*;
+import com.intellij.psi.impl.PsiDocumentManagerImpl;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+
+/**
+ * Created by IntelliJ IDEA.
+ * Author: Alexey.Ivanov
+ * Date:   23.04.2010
+ * Time:   16:44:57
+ */
+public class StatementMover extends LineMover {
+  private static final Logger LOG = Logger.getInstance("#com.jetbrains.python.codeInsight.editorActions.moveUpDown.StatementMover");
+
+  private static PsiElement[] findStatementsInRange(PsiFile file, int startOffset, int endOffset) {
+    PsiElement element1 = file.findElementAt(startOffset);
+    PsiElement element2 = file.findElementAt(endOffset - 1);
+    if (element1 instanceof PsiWhiteSpace) {
+      startOffset = element1.getTextRange().getEndOffset();
+      element1 = file.findElementAt(startOffset);
+    }
+    if (element2 instanceof PsiWhiteSpace) {
+      endOffset = element2.getTextRange().getStartOffset();
+      element2 = file.findElementAt(endOffset - 1);
+    }
+    if (element1 == null || element2 == null) {
+      return PsiElement.EMPTY_ARRAY;
+    }
+
+    PsiElement parent = PsiTreeUtil.findCommonParent(element1, element2);
+    if (parent == null) {
+      return PsiElement.EMPTY_ARRAY;
+    }
+
+    while (true) {
+      if (parent instanceof PyStatement) {
+        parent = parent.getParent();
+        break;
+      }
+      if (parent instanceof PyStatementList) {
+        break;
+      }
+      if (parent == null || parent instanceof PsiFile) {
+        return PsiElement.EMPTY_ARRAY;
+      }
+      parent = parent.getParent();
+    }
+
+    if (!parent.equals(element1)) {
+      while (!parent.equals(element1.getParent())) {
+        element1 = element1.getParent();
+      }
+    }
+
+    if (!parent.equals(element2)) {
+      while (!parent.equals(element2.getParent())) {
+        element2 = element2.getParent();
+      }
+    }
+
+    PsiElement[] children = parent.getChildren();
+    ArrayList<PsiElement> array = new ArrayList<PsiElement>();
+    boolean flag = false;
+    for (PsiElement child : children) {
+      if (child.equals(element1)) {
+        flag = true;
+      }
+      if (flag && !(child instanceof PsiWhiteSpace)) {
+        array.add(child);
+      }
+      if (child.equals(element2)) {
+        break;
+      }
+    }
+
+    for (PsiElement element : array) {
+      if (!(element instanceof PyStatement || element instanceof PsiWhiteSpace || element instanceof PsiComment)) {
+        return PsiElement.EMPTY_ARRAY;
+      }
+    }
+    return array.toArray(new PsiElement[array.size()]);
+  }
+
+  private static boolean isNotValidStatementRange(Pair<PsiElement, PsiElement> range) {
+    return range == null ||
+           range.first == null ||
+           range.second == null ||
+           range.first.getParent() != range.second.getParent();
+  }
+
+  @Nullable
+  private static LineRange expandLineRange(LineRange range, Editor editor, PsiFile file) {
+    final SelectionModel selectionModel = editor.getSelectionModel();
+    Pair<PsiElement, PsiElement> psiRange;
+    if (selectionModel.hasSelection()) {
+      final int startOffset = selectionModel.getSelectionStart();
+      final int endOffset = selectionModel.getSelectionEnd();
+      final PsiElement[] psiElements = findStatementsInRange(file, startOffset, endOffset);
+      if (psiElements.length == 0) {
+        return null;
+      }
+      psiRange = new Pair<PsiElement, PsiElement>(psiElements[0], psiElements[psiElements.length - 1]);
+    }
+    else {
+      final int offset = editor.getCaretModel().getOffset();
+      PsiElement element = file.findElementAt(offset);
+      psiRange = new Pair<PsiElement, PsiElement>(element, element);
+    }
+
+    psiRange = new Pair<PsiElement, PsiElement>(PsiTreeUtil.getNonStrictParentOfType(psiRange.getFirst(), PyStatement.class),
+                                                PsiTreeUtil.getNonStrictParentOfType(psiRange.getSecond(), PyStatement.class));
+    if (psiRange.getFirst() == null || psiRange.getSecond() == null) {
+      return null;
+    }
+
+    final PsiElement parent = PsiTreeUtil.findCommonParent(psiRange.getFirst(), psiRange.getSecond());
+    final Pair<PsiElement, PsiElement> elementRange = getElementRange(parent, psiRange.getFirst(), psiRange.getSecond());
+    if (isNotValidStatementRange(elementRange)) {
+      return null;
+    }
+
+    if (elementRange.getFirst() == elementRange.getSecond() && elementRange.getFirst() instanceof PyPassStatement) {
+      return null;
+    }
+
+    final int endOffset = elementRange.getSecond().getTextRange().getEndOffset();
+    final Document document = editor.getDocument();
+    if (endOffset > document.getTextLength()) {
+      LOG.assertTrue(!PsiDocumentManager.getInstance(file.getProject()).isUncommited(document));
+      LOG.assertTrue(PsiDocumentManagerImpl.checkConsistency(file, document));
+    }
+
+    int endLine;
+    if (endOffset == document.getTextLength()) {
+      endLine = document.getLineCount();
+    }
+    else {
+      endLine = editor.offsetToLogicalPosition(endOffset).line + 1;
+      endLine = Math.min(endLine, document.getLineCount());
+    }
+    endLine = Math.max(endLine, range.endLine);
+    final int startLine = Math.min(range.startLine, editor.offsetToLogicalPosition(elementRange.getFirst().getTextOffset()).line);
+    return new LineRange(startLine, endLine);
+  }
+
+  private @Nullable PyStatementList myStatementListToAddPass;
+  private @Nullable PyStatementList myStatementListToRemovePass;
+  private @NotNull PsiElement[] myElementsToDecreaseIndent;
+  private @NotNull PsiElement[] myElementsToIncreaseIndent;
+
+  @Override
+  public boolean checkAvailable(@NotNull Editor editor,
+                                @NotNull PsiFile file,
+                                @NotNull MoveInfo info,
+                                boolean down) {
+    myStatementListToAddPass = null;
+    myStatementListToRemovePass = null;
+    myElementsToDecreaseIndent = PsiElement.EMPTY_ARRAY;
+    myElementsToIncreaseIndent = PsiElement.EMPTY_ARRAY;
+    if (!super.checkAvailable(editor, file, info, down)) {
+      return false;
+    }
+    info.indentSource = true;
+    final LineRange range = expandLineRange(info.toMove, editor, file);
+    if (range == null) {
+      return false;
+    }
+
+    info.toMove = range;
+    final int startOffset = editor.logicalPositionToOffset(new LogicalPosition(range.startLine, 0));
+    final int endOffset = editor.logicalPositionToOffset(new LogicalPosition(range.endLine, 0));
+    final PsiElement[] statements = findStatementsInRange(file, startOffset, endOffset);
+    final int length = statements.length;
+    if (length == 0) {
+      return false;
+    }
+
+    range.firstElement = statements[0];
+    range.lastElement = statements[length - 1];
+    final Document document = editor.getDocument();
+    PyStatement statement;
+    assert info.toMove2.startLine + 1 == info.toMove2.endLine;
+    if (down) {
+      statement = PsiTreeUtil.getNextSiblingOfType(range.lastElement, PyStatement.class);
+    }
+    else {
+      statement = PsiTreeUtil.getPrevSiblingOfType(range.firstElement, PyStatement.class);
+    }
+
+    if (statement == null) {
+      final PyStatementPart parentStatementPart =
+        PsiTreeUtil.getParentOfType(PsiTreeUtil.findCommonParent(range.firstElement, range.lastElement), PyStatementPart.class, false);
+      if (parentStatementPart == null) {
+        info.toMove2 = null;
+      }
+      else {  //we are in statement part
+
+        final PyStatementList statementList = parentStatementPart.getStatementList();
+        assert statementList != null;
+        if (down) {
+          final PyStatementPart nextStatementPart = PsiTreeUtil.getNextSiblingOfType(statementList.getParent(), PyStatementPart.class);
+          if (nextStatementPart != null) {
+            info.toMove2 = new LineRange(range.endLine, range.endLine + 1);
+            myStatementListToRemovePass = nextStatementPart.getStatementList();
+          }
+          else {
+            final PyStatement parentStatement = PsiTreeUtil.getParentOfType(statementList, PyStatement.class);
+            if (parentStatement == null) {
+              return false;
+            }
+            final PyStatement nextStatement = PsiTreeUtil.getNextSiblingOfType(parentStatement, PyStatement.class);
+            if (nextStatement == null) {
+              return false;
+            }
+            final int startLine = editor.offsetToLogicalPosition(parentStatement.getTextRange().getEndOffset()).line;
+            final int endLine = editor.offsetToLogicalPosition(nextStatement.getTextRange().getEndOffset()).line;
+            info.toMove2 = new LineRange(startLine + 1, endLine + 1);
+
+            myElementsToDecreaseIndent = statements;
+          }
+        }
+        else {
+          final PyStatementPart prevStatementPart = PsiTreeUtil.getPrevSiblingOfType(statementList.getParent(), PyStatementPart.class);
+          if (prevStatementPart != null) {
+            myStatementListToRemovePass = prevStatementPart.getStatementList();
+          }
+          else {
+            myElementsToDecreaseIndent = statements;
+            info.toMove2 = new LineRange(range.startLine - 1, range.startLine);
+          }
+        }
+        if (Arrays.equals(statementList.getStatements(), statements)) {
+          myStatementListToAddPass = statementList;
+        }
+      }
+      return true;
+    }
+
+    info.toMove2 = new LineRange(statement, statement, document);
+
+    final PyStatementPart[] statementParts = PsiTreeUtil.getChildrenOfType(statement, PyStatementPart.class);
+    // next/previous statement has a statement parts
+    if (statementParts != null) {
+      // move inside statement part
+      if (down) {
+        final PyStatementPart statementPart = statementParts[0];
+        final int lineNumber = document.getLineNumber(statementPart.getTextRange().getStartOffset());
+        info.toMove2 = new LineRange(lineNumber, lineNumber + 1);
+        myStatementListToRemovePass = statementPart.getStatementList();
+        myElementsToIncreaseIndent = statements;
+      }
+      else {
+        final PyStatementPart statementPart = statementParts[statementParts.length - 1];
+        final int lineNumber = document.getLineNumber(statementPart.getTextRange().getEndOffset());
+        info.toMove2 = new LineRange(lineNumber, lineNumber + 1);
+        myStatementListToRemovePass = statementPart.getStatementList();
+        myElementsToIncreaseIndent = statements;
+      }
+    }
+    return true;
+  }
+
+  private void decreaseIndent(final Editor editor) {
+    final Document document = editor.getDocument();
+    for (PsiElement statement : myElementsToDecreaseIndent) {
+      final int startOffset = statement.getTextRange().getStartOffset() - 1;
+      PsiElement element = statement.getContainingFile().findElementAt(startOffset);
+      assert element instanceof PsiWhiteSpace;
+      final String text = element.getText();
+      String[] lines = text.split("\n");
+      if (lines.length == 0) {
+        continue;
+      }
+      final int indent = lines[lines.length - 1].length();
+      final int startLine = editor.offsetToLogicalPosition(startOffset).line;
+      final int endLine = editor.offsetToLogicalPosition(statement.getTextRange().getEndOffset()).line;
+      for (int line = startLine; line <= endLine; ++line) {
+        if (indent >= 4) {
+          final int lineStartOffset = document.getLineStartOffset(line);
+          document.deleteString(lineStartOffset, lineStartOffset + 4);
+        }
+      }
+    }
+  }
+
+  private void increaseIndent(final Editor editor) {
+    final Document document = editor.getDocument();
+    for (PsiElement statement : myElementsToIncreaseIndent) {
+      final int startLine = editor.offsetToLogicalPosition(statement.getTextRange().getStartOffset()).line;
+      final int endLine = editor.offsetToLogicalPosition(statement.getTextRange().getEndOffset()).line;
+      for (int line = startLine; line <= endLine; ++line) {
+        final int offset = document.getLineStartOffset(line);
+        document.insertString(offset, "    ");
+      }
+    }
+  }
+
+  @Override
+  public void beforeMove(@NotNull Editor editor, @NotNull MoveInfo info, boolean down) {
+    super.beforeMove(editor, info, down);
+    if (myStatementListToAddPass != null) {
+      final PyPassStatement passStatement =
+        PyElementGenerator.getInstance(editor.getProject()).createFromText(PyPassStatement.class, "pass");
+      myStatementListToAddPass.add(passStatement);
+      CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(myStatementListToAddPass);
+      if (down) {
+        info.toMove2 = new LineRange(info.toMove2.startLine, info.toMove2.endLine + 1);
+      }
+    }
+    decreaseIndent(editor);
+    increaseIndent(editor);
+  }
+
+  @Override
+  public void afterMove(@NotNull Editor editor,
+                        @NotNull PsiFile file,
+                        @NotNull MoveInfo info,
+                        boolean down) {
+    super.afterMove(editor, file, info, down);
+    if (myStatementListToRemovePass != null) {
+      PyPsiUtils.removeRedundantPass(myStatementListToRemovePass);
+      CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(myStatementListToRemovePass);
+    }
+  }
+}

python/testData/mover/betweenStatementParts.py

@@ -0,0 +1,8 @@
+class A:
+    def f(self, a, b):
+        if a:
+            b = 1
+        elif b:
+            a = <caret>1
+        else:
+            print "a"

python/testData/mover/betweenStatementParts_afterDown.py

@@ -0,0 +1,9 @@
+class A:
+    def f(self, a, b):
+        if a:
+            b = 1
+        elif b:
+            pass
+        else:
+            a = 1
+            print "a"

python/testData/mover/betweenStatementParts_afterUp.py

@@ -0,0 +1,9 @@
+class A:
+    def f(self, a, b):
+        if a:
+            b = 1
+            a = 1
+        elif b:
+            pass
+        else:
+            print "a"

python/testData/mover/functions.py

@@ -0,0 +1,9 @@
+class A:
+    def a(self):
+        doSomething()
+
+    def b<caret>(self):
+        doSomething()
+
+    def c(self):
+        doSomething()

python/testData/mover/functions_afterDown.py

@@ -0,0 +1,9 @@
+class A:
+    def a(self):
+        doSomething()
+
+    def c(self):
+        doSomething()
+
+    def b(self):
+        doSomething()

python/testData/mover/functions_afterUp.py

@@ -0,0 +1,9 @@
+class A:
+    def b(self):
+        doSomething()
+
+    def a(self):
+        doSomething()
+
+    def c(self):
+        doSomething()

python/testData/mover/insideStatement.py

@@ -0,0 +1,8 @@
+def foo(a, b):
+    if a:
+        pass
+    else:
+        pass
+    a = <caret>3
+    if b:
+        pass

python/testData/mover/insideStatement_afterDown.py

@@ -0,0 +1,7 @@
+def foo(a, b):
+    if a:
+        pass
+    else:
+        pass
+    if b:
+        a = 3

python/testData/mover/insideStatement_afterUp.py

@@ -0,0 +1,7 @@
+def foo(a, b):
+    if a:
+        pass
+    else:
+        a = 3
+    if b:
+        pass

python/testData/mover/moveStatement.py

@@ -0,0 +1,12 @@
+class A:
+    def f(self, a, b):
+        while a:
+            a = a / 2
+        if a:
+            b = 1
+        elif<caret> b:
+            a = 1
+        else:
+            print "a"
+        for i in range(1, 10):
+            print i

python/testData/mover/moveStatement_afterDown.py

@@ -0,0 +1,12 @@
+class A:
+    def f(self, a, b):
+        while a:
+            a = a / 2
+        for i in range(1, 10):
+            if a:
+                b = 1
+            elif b:
+                a = 1
+            else:
+                print "a"
+            print i

python/testData/mover/moveStatement_afterUp.py

@@ -0,0 +1,12 @@
+class A:
+    def f(self, a, b):
+        while a:
+            if a:
+                b = 1
+            elif b:
+                a = 1
+            else:
+                print "a"
+            a = a / 2
+        for i in range(1, 10):
+            print i

python/testData/mover/outsideStatement.py

@@ -0,0 +1,6 @@
+a = 1
+
+if a:
+    a <caret>= 3
+
+a = 4

python/testData/mover/outsideStatement_afterDown.py

@@ -0,0 +1,7 @@
+a = 1
+
+if a:
+    pass
+
+a = 4
+a = 3

python/testData/mover/outsideStatement_afterUp.py

@@ -0,0 +1,7 @@
+a = 1
+
+a = 3
+if a:
+    pass
+
+a = 4

python/testData/mover/selection.py

@@ -0,0 +1,11 @@
+def f(a, b, c):
+    if b:
+        pass
+    elif a:
+        c = <selection>1
+        if b:
+            a = 3
+        else:
+            b<caret></selection> = 4
+    else:
+        pass

python/testData/mover/selection_afterDown.py

@@ -0,0 +1,11 @@
+def f(a, b, c):
+    if b:
+        pass
+    elif a:
+        pass
+    else:
+        c = <selection>1
+        if b:
+            a = 3
+        else:
+            b</selection> = 4

python/testData/mover/selection_afterUp.py

@@ -0,0 +1,11 @@
+def f(a, b, c):
+    if b:
+        c = <selection>1
+        if b:
+            a = 3
+        else:
+            b</selection> = 4
+    elif a:
+        pass
+    else:
+        pass

python/testData/mover/simple.py

@@ -0,0 +1,2 @@
+a =<caret> 1
+b = 2

python/testData/mover/simple_afterDown.py

@@ -0,0 +1,2 @@
+b = 2
+a = 1

python/testData/mover/simple_afterUp.py

@@ -0,0 +1,2 @@
+a = 1
+b = 2

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

@@ -0,0 +1,63 @@
+package com.jetbrains.python;
+
+import com.intellij.codeInsight.editorActions.moveUpDown.MoveStatementDownAction;
+import com.intellij.codeInsight.editorActions.moveUpDown.MoveStatementUpAction;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.actionSystem.EditorActionHandler;
+import com.intellij.openapi.fileEditor.FileDocumentManager;
+import com.jetbrains.python.fixtures.PyLightFixtureTestCase;
+
+/**
+ * Created by IntelliJ IDEA.
+ * Author: Alexey.Ivanov
+ * Date:   05.05.2010
+ * Time:   18:20:20
+ */
+public class PyStatementMoverTest extends PyLightFixtureTestCase {
+  private void doTest() throws Exception {
+    final String testName = getTestName(true);
+    myFixture.configureByFile("mover/" + testName + ".py");
+    performAction(new MoveStatementUpAction().getHandler());
+    myFixture.checkResultByFile("mover/" + testName + "_afterUp.py", true);
+    
+    FileDocumentManager.getInstance().reloadFromDisk(myFixture.getDocument(myFixture.getFile()));
+    myFixture.configureByFile("mover/" + getTestName(true) + ".py");
+    performAction(new MoveStatementDownAction().getHandler());
+    myFixture.checkResultByFile("mover/" + testName + "_afterDown.py", true);
+  }
+
+  private void performAction(EditorActionHandler handler) {
+    final Editor editor = myFixture.getEditor();
+    if (handler.isEnabled(editor, null)) {
+      handler.execute(editor, null);
+    }
+  }
+
+  public void testSimple() throws Exception {
+    doTest();
+  }
+
+  public void testOutsideStatement() throws Exception {
+    doTest();
+  }
+
+  public void testInsideStatement() throws Exception {
+    doTest();
+  }
+
+  public void testFunctions() throws Exception {
+    doTest();
+  }
+
+  public void testBetweenStatementParts() throws Exception {
+    doTest();
+  }
+
+  public void testMoveStatement() throws Exception {
+    doTest();
+  }
+
+  public void testSelection() throws Exception {
+    doTest();
+  }
+}