PY-349

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -60,6 +60,23 @@ public class PyUtil {
   }
 
 
+  @NotNull
+  public static Set<PsiElement> getComments(PsiElement start) {
+    final Set<PsiElement> comments = new HashSet<PsiElement>();
+    PsiElement seeker = start.getPrevSibling();
+    if (seeker == null) seeker = start.getParent().getPrevSibling();
+    while (seeker instanceof PsiWhiteSpace || seeker instanceof PsiComment) {
+      if (seeker instanceof PsiComment) {
+        comments.add(seeker);
+      }
+      seeker = seeker.getPrevSibling();
+    }
+    if (seeker instanceof PyExpressionStatement && seeker.getFirstChild() instanceof PyStringLiteralExpression) {
+      comments.add(seeker);
+    }
+    return comments;
+  }
+
   @Nullable
   public static PsiElement getFirstNonCommentAfter(PsiElement start) {
     PsiElement seeker = start;

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

@@ -91,6 +91,16 @@ public class PyClassRefactoringUtil {
     if (methods.size() == 0) return;
     final PyElement[] elements = methods.toArray(new PyElement[methods.size()]);
     addMethods(superClass, elements, true);
+    removeMethodsWithComments(elements);
+  }
+
+  private static void removeMethodsWithComments(PyElement[] elements) {
+    for (PyElement element : elements) {
+      final Set<PsiElement> comments = PyUtil.getComments(element);
+      if (comments.size() > 0) {
+        PyPsiUtils.removeElements(comments.toArray(new PsiElement[comments.size()]));
+      }
+    }
     PyPsiUtils.removeElements(elements);
   }
 
@@ -110,13 +120,18 @@ public class PyClassRefactoringUtil {
 
     final PyClass newClass = PyElementGenerator.getInstance(project).createFromText(PyClass.class, text);
     final PyStatementList statements = superClass.getStatementList();
+    final PyStatementList newStatements = newClass.getStatementList();
     if (statements.getStatements().length != 0) {
-      for (PyElement newStatement : newClass.getStatementList().getStatements()) {
-        //statements.add(PythonLanguage.getInstance().getElementGenerator().createNewLine(project));
-        statements.add(newStatement);
+      for (PyElement newStatement : newStatements.getStatements()) {
+        if (newStatement instanceof PyExpressionStatement && newStatement.getFirstChild() instanceof PyStringLiteralExpression) continue;
+        final PsiElement anchor = statements.add(newStatement);
+        final Set<PsiElement> comments = PyUtil.getComments(newStatement);
+        for (PsiElement comment : comments) {
+          statements.addBefore(comment, anchor);
+        }
       }
     } else {
-      statements.replace(newClass.getStatementList());
+      statements.replace(newStatements);
     }
   }
 
@@ -134,6 +149,10 @@ public class PyClassRefactoringUtil {
     for (PyElement element : elements) {
       final String name = element.getName();
       if (name != null && (up || superClass.findMethodByName(name, false) == null)) {
+        final Set<PsiElement> comments = PyUtil.getComments(element);
+        for (PsiElement comment : comments) {
+          builder.append(white).append(comment.getText());
+        }
         builder.append(white).append(element.getText()).append("\n");
         hasChanges = true;
       }

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

@@ -0,0 +1,10 @@
+class Foo:
+    def foo(self):
+        print("a")
+
+    # this is boo
+    def boo(self):
+        print "rrrrr"
+
+class Boo(Foo):
+    pass

python/testData/refactoring/pullup/withComments.py

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

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

@@ -0,0 +1,14 @@
+class Foo:
+    def foo(self):
+        print("a")
+
+    '''
+    this is boo
+    very long boo
+    '''
+
+    def boo(self):
+        print "rrrrr"
+
+class Boo(Foo):
+    pass

python/testData/refactoring/pullup/withMultilineComments.py

@@ -0,0 +1,11 @@
+class Foo:
+    def foo(self):
+        print("a")
+
+class Boo(Foo):
+    '''
+    this is boo
+    very long boo
+    '''
+    def boo(self):
+        print "rrrrr"

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

@@ -10,19 +10,27 @@ import java.util.Collections;
  * @author Dennis.Ushakov
  */
 public class PyPullUpTest extends PyClassRefactoringTest {
-  public void testSimple() throws Exception {
+  public void testSimple() {
     doHelperTest("Boo", ".boo", "Foo");
   }
 
-  public void testSuperclass() throws Exception {
+  public void testSuperclass() {
     doHelperTest("Boo", "Foo", "Zope");
   }
 
-  public void testExistingsuperclass() throws Exception {
+  public void testExistingsuperclass() {
     doHelperTest("Boo", "Foo", "Zope");
   }
 
-  private void doHelperTest(final String className, final String memberName, final String superClassName) throws Exception {
+  public void testWithComments() {
+    doHelperTest("Boo", ".boo", "Foo");
+  }
+
+  public void testWithMultilineComments() {
+    doHelperTest("Boo", ".boo", "Foo");
+  }
+
+  private void doHelperTest(final String className, final String memberName, final String superClassName) {
     String baseName = "/refactoring/pullup/" + getTestName(true);
     myFixture.configureByFile(baseName + ".py");
     final PyClass clazz = findClass(className);