pull up snake

python/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -4,6 +4,7 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.Pair;
@@ -19,6 +20,7 @@ import org.jetbrains.annotations.Nullable;
 public class PyPsiUtils {
   public static final Key<Pair<PsiElement, TextRange>> SELECTION_BREAKS_AST_NODE =
     new Key<Pair<PsiElement, TextRange>>("python.selection.breaks.ast.node");
+  private static final Logger LOG = Logger.getInstance(PyPsiUtils.class.getName());
 
   private PyPsiUtils() {
   }
@@ -72,4 +74,35 @@ public class PyPsiUtils {
       oldExpression.replace(newExpression);
     }
   }
+
+  public static void addToEnd(@NotNull final PsiElement psiElement, @NotNull final PsiElement... newElements) {
+    final ASTNode psiNode = psiElement.getNode();
+    LOG.assertTrue(psiNode != null);
+    for (PsiElement newElement : newElements) {
+      //noinspection ConstantConditions
+      psiNode.addChild(newElement.getNode());
+    }
+  }
+
+  public static void addBeforeInParent(@NotNull final PsiElement anchor, @NotNull final PsiElement... newElements) {
+    final PsiElement psiParent = anchor.getParent();
+    LOG.assertTrue(psiParent != null);
+    final ASTNode parentNode = psiParent.getNode();
+    final ASTNode anchorNode = anchor.getNode();
+    LOG.assertTrue(parentNode != null);
+    LOG.assertTrue(anchorNode != null);
+    for (PsiElement newElement : newElements) {
+      //noinspection ConstantConditions
+      parentNode.addChild(newElement.getNode(), anchorNode);
+    }
+  }
+
+  public static void removeElements(@NotNull final PsiElement... elements) {
+    final ASTNode parentNode = elements[0].getParent().getNode();
+    LOG.assertTrue(parentNode != null);
+    for (PsiElement element : elements) {
+      //noinspection ConstantConditions
+      parentNode.removeChild(element.getNode());
+    }
+  }
 }

python/src/com/jetbrains/python/refactoring/classes/PyMemberInfo.java

@@ -2,6 +2,7 @@ package com.jetbrains.python.refactoring.classes;
 
 import com.intellij.refactoring.classMembers.MemberInfoBase;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.refactoring.classes.ui.PyClassCellRenderer;
 
 /**
  * @author Dennis.Ushakov
@@ -21,6 +22,8 @@ public class PyMemberInfo extends MemberInfoBase<PyElement> {
           overrides = true;
         }
       }
+    } else if (member instanceof PyClass) {
+      displayName = PyClassCellRenderer.getClassText((PyClass)member);
     }
   }
 
@@ -37,4 +40,17 @@ public class PyMemberInfo extends MemberInfoBase<PyElement> {
     builder.append(")");
     return builder.toString();
   }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj instanceof PyMemberInfo) {
+      return getMember().equals(((PyMemberInfo)obj).getMember());
+    }
+    return false;
+  }
+
+  @Override
+  public int hashCode() {
+    return getMember().hashCode();
+  }
 }

python/src/com/jetbrains/python/refactoring/classes/PyMemberInfoStorage.java

@@ -55,6 +55,9 @@ public class PyMemberInfoStorage extends AbstractMemberInfoStorage<PyElement, Py
     for (PyFunction function : aClass.getMethods()) {
       temp.add(new PyMemberInfo(function));
     }
+    for (PyClass pyClass : aClass.getSuperClasses()) {
+      temp.add(new PyMemberInfo(pyClass));
+    }
   }
 
   @Override

python/src/com/jetbrains/python/refactoring/classes/pullUp/PullUpConflictsUtil.java

@@ -27,6 +27,14 @@ public class PullUpConflictsUtil {
         if (name == null) continue;
         final PyFunction superClassMethod = superClass.findMethodByName(name, false);
         isConflict = superClassMethod != null;
+      } else if (member instanceof PyClass) {
+        final PyClass clazz = (PyClass)member;
+        for (PyClass aClass : superClass.getSuperClasses()) {
+          if (aClass == clazz) {
+            conflictsList.putValue(superClass,
+                                   RefactoringUIUtil.getDescription(superClass, false) + " already extends " + RefactoringUIUtil.getDescription(clazz, false));
+          }
+        }
       }
 
       if (isConflict) {

python/src/com/jetbrains/python/refactoring/classes/pullUp/PyPullUpDialog.java

@@ -139,8 +139,9 @@ public class PyPullUpDialog extends UpDirectedMembersMovingDialog {
 
     public boolean isMemberEnabled(PyMemberInfo member) {
       PyClass currentSuperClass = getSuperClass();
-      return currentSuperClass == null ||
-             !myStorage.getDuplicatedMemberInfos(currentSuperClass).contains(member);
+      return (currentSuperClass == null ||
+             !myStorage.getDuplicatedMemberInfos(currentSuperClass).contains(member)) &&
+             member.getMember() != currentSuperClass;
     }
 
     public boolean isAbstractEnabled(PyMemberInfo member) {

python/src/com/jetbrains/python/refactoring/classes/pullUp/PyPullUpHandler.java

@@ -2,16 +2,16 @@ package com.jetbrains.python.refactoring.classes.pullUp;
 
 import com.intellij.openapi.actionSystem.DataContext;
 import com.intellij.openapi.application.ex.ApplicationManagerEx;
-import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.refactoring.RefactoringBundle;
+import com.intellij.refactoring.lang.ElementsHandler;
 import com.intellij.refactoring.util.CommonRefactoringUtil;
+import com.intellij.util.PsiNavigateUtil;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.psi.PyClass;
-import com.jetbrains.python.psi.PyElement;
 import com.jetbrains.python.psi.PyUtil;
 import com.jetbrains.python.refactoring.classes.PyClassMembersRefactoringSupport;
 import com.jetbrains.python.refactoring.classes.PyClassRefactoringHandler;
@@ -23,9 +23,8 @@ import java.util.Collection;
 /**
  * @author: Dennis.Ushakov
  */
-public class PyPullUpHandler extends PyClassRefactoringHandler {
+public class PyPullUpHandler extends PyClassRefactoringHandler implements ElementsHandler {
   public static final String REFACTORING_NAME = PyBundle.message("refactoring.pull.up.dialog.title");
-  private static final Logger LOG = Logger.getInstance("com.jetbrains.python.refactoring.classes.pullUp");
 
   @Override
   protected void doRefactor(Project project, PsiElement element1, PsiElement element2, Editor editor, PsiFile file, DataContext dataContext) {
@@ -53,8 +52,8 @@ public class PyPullUpHandler extends PyClassRefactoringHandler {
     }
   }
 
-  private void pullUpWithHelper(PyClass clazz, Collection<PyMemberInfo> selectedMemberInfos, PyClass superClass) {
-
+  private static void pullUpWithHelper(PyClass clazz, Collection<PyMemberInfo> selectedMemberInfos, PyClass superClass) {
+    PsiNavigateUtil.navigate(PyPullUpHelper.pullUp(clazz, selectedMemberInfos, superClass));
   }
 
   @Override
@@ -66,4 +65,8 @@ public class PyPullUpHandler extends PyClassRefactoringHandler {
   protected String getHelpId() {
     return "refactoring.pullMembersUp";
   }
+
+  public boolean isEnabledOnElements(PsiElement[] elements) {
+    return elements.length == 1 && elements[0] instanceof PyClass;
+  }
 }

python/src/com/jetbrains/python/refactoring/classes/pullUp/PyPullUpHelper.java

@@ -0,0 +1,110 @@
+package com.jetbrains.python.refactoring.classes.pullUp;
+
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.command.CommandProcessor;
+import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiFileFactory;
+import com.intellij.refactoring.RefactoringBundle;
+import com.jetbrains.python.PythonFileType;
+import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyFunction;
+import com.jetbrains.python.psi.PyUtil;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
+import com.jetbrains.python.refactoring.classes.PyMemberInfo;
+
+import java.util.*;
+
+/**
+ * @author Dennis.Ushakov
+ */
+public class PyPullUpHelper {
+  private static final Logger LOG = Logger.getInstance(PyPullUpHelper.class.getName());
+  private PyPullUpHelper() {}
+
+  public static PyElement pullUp(final PyClass clazz, final Collection<PyMemberInfo> selectedMemberInfos, final PyClass superClass) {
+    final Set<String> superClasses = new HashSet<String>();
+    final List<PyFunction> methods = new ArrayList<PyFunction>();
+    for (PyMemberInfo member : selectedMemberInfos) {
+      final PyElement element = member.getMember();
+      if (element instanceof PyFunction) methods.add((PyFunction)element);
+      else if (element instanceof PyClass) superClasses.add(element.getName());
+      else LOG.error("unmatched member class " + element.getClass());
+    }
+    CommandProcessor.getInstance().executeCommand(clazz.getProject(), new Runnable() {
+      public void run() {
+        ApplicationManager.getApplication().runWriteAction(new Runnable() {
+          public void run() {
+            // move methods
+            moveElements(methods, clazz, superClass);
+
+            // move superclasses declarations
+            moveSuperclasses(clazz, superClasses, superClass);
+          }
+        });
+      }
+    }, RefactoringBundle.message("pull.members.up.title"), null);
+
+    return superClass;
+  }
+
+  private static void moveSuperclasses(PyClass clazz, Set<String> superClasses, PyClass superClass) {
+    if (superClasses.size() == 0) return;
+    final Project project = clazz.getProject();
+    final List<PsiElement> toAdd = new ArrayList<PsiElement>();
+    PsiElement[] elements = clazz.getSuperClassExpressions();
+    for (PsiElement element : elements) {
+      if (superClasses.contains(element.getText())) {
+        toAdd.add(element);
+        PyUtil.removeListNode(element);
+      }
+    }
+    elements = superClass.getSuperClassExpressions();
+
+    if (elements.length > 0) {
+      PsiElement parent = elements[elements.length - 1].getParent();
+      for (PsiElement element : toAdd) {
+        PyUtil.addListNode(parent, element, parent.getLastChild().getNode(), false, true);
+      }
+    } else {
+      addSuperclasses(project, superClass, superClasses);
+    }
+  }
+
+  private static void addSuperclasses(Project project, PyClass superClass, Collection<String> superClasses) {
+    final StringBuilder builder = new StringBuilder("(");
+    for (String element : superClasses) {
+      if (builder.length() > 1) builder.append(",");
+      builder.append(element);
+    }
+    builder.append(")");
+    final PsiFile file = PsiFileFactory.getInstance(project).createFileFromText(superClass.getName() + "temp", PythonFileType.INSTANCE, builder.toString());
+    final PsiElement expression = file.getFirstChild().getFirstChild();
+    PsiElement colon = superClass.getFirstChild();
+    while (colon != null && !colon.getText().equals(":")) {
+      colon = colon.getNextSibling();
+    }
+    PyPsiUtils.addBeforeInParent(colon, expression);
+  }
+
+  private static void moveElements(List<PyFunction> methods, PyClass clazz, PyClass superClass) {
+    if (methods.size() == 0) return;
+    final Project project = clazz.getProject();
+    final PsiElement[] elements = methods.toArray(new PsiElement[methods.size()]);
+    PyPsiUtils.removeElements(elements);
+    final StringBuilder builder = new StringBuilder();
+    for (PsiElement element : elements) {
+      builder.append(element.getText()).append("\n");
+    }
+    final PsiFile file = PsiFileFactory.getInstance(project).createFileFromText(clazz.getName() + "temp", PythonFileType.INSTANCE, builder.toString());
+
+    if (superClass.getMethods().length != 0) {
+      PyPsiUtils.addBeforeInParent(superClass.getMethods()[0], file.getChildren());
+    } else {
+      PyPsiUtils.addToEnd(superClass, file.getChildren());
+    }
+  }
+}

python/testData/refactoring/pullup/existingsuperclass.after.py

@@ -0,0 +1,11 @@
+class Foo:
+    def foo(self):
+        print("a")
+
+class Zope(Hand,Foo):
+    def _mine(self):
+       print "zope"
+
+class Boo():
+    def boo(self):
+        print "rrrrr"

python/testData/refactoring/pullup/existingsuperclass.py

@@ -0,0 +1,11 @@
+class Foo:
+    def foo(self):
+        print("a")
+
+class Zope(Hand):
+    def _mine(self):
+       print "zope"
+
+class Boo(Foo):
+    def boo(self):
+        print "rrrrr"

python/testData/refactoring/pullup/simple.after.py

@@ -0,0 +1,8 @@
+class Foo:
+    def boo(self):
+        print "rrrrr"
+    def foo(self):
+        print("a")
+
+class Boo(Foo):
+    

python/testData/refactoring/pullup/simple.py

@@ -0,0 +1,7 @@
+class Foo:
+    def foo(self):
+        print("a")
+
+class Boo(Foo):
+    def boo(self):
+        print "rrrrr"

python/testData/refactoring/pullup/superclass.after.py

@@ -0,0 +1,11 @@
+class Foo:
+    def foo(self):
+        print("a")
+
+class Zope(Foo):
+    def _mine(self):
+       print "zope"
+
+class Boo():
+    def boo(self):
+        print "rrrrr"

python/testData/refactoring/pullup/superclass.py

@@ -0,0 +1,11 @@
+class Foo:
+    def foo(self):
+        print("a")
+
+class Zope:
+    def _mine(self):
+       print "zope"
+
+class Boo(Foo):
+    def boo(self):
+        print "rrrrr"

python/testSrc/com/jetbrains/python/refactoring/classes/PyPullUpTest.java

@@ -0,0 +1,67 @@
+package com.jetbrains.python.refactoring.classes;
+
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.search.ProjectScope;
+import com.intellij.psi.stubs.StubIndex;
+import com.jetbrains.python.PythonTestUtil;
+import com.jetbrains.python.fixtures.PyLightFixtureTestCase;
+import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyFunction;
+import com.jetbrains.python.psi.stubs.PyClassNameIndex;
+import com.jetbrains.python.refactoring.classes.pullUp.PyPullUpHelper;
+
+import java.util.Collection;
+import java.util.Collections;
+
+/**
+ * @author Dennis.Ushakov
+ */
+public class PyPullUpTest extends PyLightFixtureTestCase {
+  public void testSimple() throws Exception {
+    doHelperTest("Boo", ".boo", "Foo");
+  }
+
+  public void testSuperclass() throws Exception {
+    doHelperTest("Boo", "Foo", "Zope");
+  }
+
+  public void testExistingsuperclass() throws Exception {
+    doHelperTest("Boo", "Foo", "Zope");
+  }
+
+  private void doHelperTest(final String className, final String memberName, final String superClassName) throws Exception {
+    String baseName = "/" + getTestName(true);
+    myFixture.configureByFile(baseName + ".py");
+    final PyClass clazz = findClass(className);
+    final PyElement member = findMember(className, memberName);
+    final PyClass superClass = findClass(superClassName);
+    PyPullUpHelper.pullUp(clazz, Collections.singleton(new PyMemberInfo(member)), superClass);
+    myFixture.checkResultByFile(baseName + ".after.py");
+  }
+
+  private PyElement findMember(String className, String memberName) {
+    if (!memberName.contains(".")) return findClass(memberName);
+    return findMethod(className, memberName.substring(1));
+  }
+
+  private PyFunction findMethod(final String className, final String name) {
+    final PyClass clazz = findClass(className);
+    final PyFunction method = clazz.findMethodByName(name, false);
+    assertNotNull(method);
+    return method;
+  }
+
+  private PyClass findClass(final String name) {
+    final Project project = myFixture.getProject();
+    final Collection<PyClass> classes = StubIndex.getInstance().get(PyClassNameIndex.KEY, name, project,
+                                                                    ProjectScope.getProjectScope(project));
+    assertEquals(1, classes.size());
+    return classes.iterator().next();
+  }
+
+  @Override
+  protected String getTestDataPath() {
+    return PythonTestUtil.getTestDataPath() + "/refactoring/pullup/";
+  }
+}