diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index 5d7a9a56ab55..cda59b38513b 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/python/src/com/jetbrains/python/psi/PyUtil.java @@ -60,6 +60,23 @@ public class PyUtil { } + @NotNull + public static Set getComments(PsiElement start) { + final Set comments = new HashSet(); + 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; diff --git a/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java b/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java index ea1493cd0219..7c544800821e 100644 --- a/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java +++ b/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 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 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 comments = PyUtil.getComments(element); + for (PsiElement comment : comments) { + builder.append(white).append(comment.getText()); + } builder.append(white).append(element.getText()).append("\n"); hasChanges = true; } diff --git a/python/testData/refactoring/pullup/withComments.after.py b/python/testData/refactoring/pullup/withComments.after.py new file mode 100644 index 000000000000..5ddd8d4241f3 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/python/testData/refactoring/pullup/withComments.py b/python/testData/refactoring/pullup/withComments.py new file mode 100644 index 000000000000..32edbc391bf2 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/python/testData/refactoring/pullup/withMultilineComments.after.py b/python/testData/refactoring/pullup/withMultilineComments.after.py new file mode 100644 index 000000000000..e1a0f1a49780 --- /dev/null +++ b/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 diff --git a/python/testData/refactoring/pullup/withMultilineComments.py b/python/testData/refactoring/pullup/withMultilineComments.py new file mode 100644 index 000000000000..f8c81f6d65bb --- /dev/null +++ b/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" diff --git a/python/testSrc/com/jetbrains/python/refactoring/classes/PyPullUpTest.java b/python/testSrc/com/jetbrains/python/refactoring/classes/PyPullUpTest.java index cb2d6b791ce1..1748f49ca095 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/classes/PyPullUpTest.java +++ b/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);