diff --git a/python/src/com/jetbrains/python/inspections/quickfix/AddCallSuperQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/AddCallSuperQuickFix.java index 6293e013f7f0..06623018bb7e 100644 --- a/python/src/com/jetbrains/python/inspections/quickfix/AddCallSuperQuickFix.java +++ b/python/src/com/jetbrains/python/inspections/quickfix/AddCallSuperQuickFix.java @@ -116,12 +116,16 @@ public class AddCallSuperQuickFix implements LocalQuickFix { } newFunction.append(superCall).append("\n\t"); boolean first = true; - for (PyStatement statement : statementList.getStatements()) { - if (first && docstring != null || statement instanceof PyPassStatement) { - first = false; - continue; + for (PsiElement child = statementList.getFirstChild(); child != null; child = child.getNextSibling()) { + if (child instanceof PyStatement) { + if (first && docstring != null || child instanceof PyPassStatement) { + first = false; + continue; + } + } + if (child instanceof PyStatement || child instanceof PsiComment) { + newFunction.append(child.getText()).append("\n\t"); } - newFunction.append(statement.getText()).append("\n\t"); } final PyElementGenerator generator = PyElementGenerator.getInstance(project); diff --git a/python/testData/inspections/AddCallSuperCommentsInFunctionBodyPreserved.py b/python/testData/inspections/AddCallSuperCommentsInFunctionBodyPreserved.py new file mode 100644 index 000000000000..9a52885bfdf5 --- /dev/null +++ b/python/testData/inspections/AddCallSuperCommentsInFunctionBodyPreserved.py @@ -0,0 +1,11 @@ +class A: + def __init__(self): + pass + + +class B(A): + def __init__(self): + # comment #1 + pass + # comment #2 + print 42 \ No newline at end of file diff --git a/python/testData/inspections/AddCallSuperCommentsInFunctionBodyPreserved_after.py b/python/testData/inspections/AddCallSuperCommentsInFunctionBodyPreserved_after.py new file mode 100644 index 000000000000..917218b68477 --- /dev/null +++ b/python/testData/inspections/AddCallSuperCommentsInFunctionBodyPreserved_after.py @@ -0,0 +1,11 @@ +class A: + def __init__(self): + pass + + +class B(A): + def __init__(self): + A.__init__(self) + # comment #1 + # comment #2 + print 42 \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java index 41a0eef005a3..b5edd40b3c6e 100644 --- a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java @@ -379,6 +379,11 @@ public class PyQuickFixTest extends PyTestCase { doInspectionTest(PyMissingConstructorInspection.class, PyBundle.message("QFIX.add.super"), true, true); } + // PY-16348 + public void testAddCallSuperCommentsInFunctionBodyPreserved() { + doInspectionTest(PyMissingConstructorInspection.class, PyBundle.message("QFIX.add.super"), true, true); + } + // PY-491, PY-13297 public void testAddEncoding() { doInspectionTest(PyMandatoryEncodingInspection.class, PyBundle.message("QFIX.add.encoding"), true, true);