diff --git a/python/src/com/jetbrains/python/inspections/quickfix/AddCallSuperQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/AddCallSuperQuickFix.java index d6568240ecf4..65dfe4537216 100644 --- a/python/src/com/jetbrains/python/inspections/quickfix/AddCallSuperQuickFix.java +++ b/python/src/com/jetbrains/python/inspections/quickfix/AddCallSuperQuickFix.java @@ -65,41 +65,36 @@ public class AddCallSuperQuickFix implements LocalQuickFix { final PyClass superClass = superClasses[0]; final PyFunction superInit = superClass.findMethodByName(PyNames.INIT, true); if (superInit == null) return; - final boolean addComma; + + final ParametersInfo origInfo = new ParametersInfo(problemFunction.getParameterList()); + final ParametersInfo superInfo = new ParametersInfo(superInit.getParameterList()); + final boolean addSelfToCall; + if (klass.isNewStyleClass()) { - addComma = false; + addSelfToCall = false; if (LanguageLevel.forElement(klass).isPy3K()) { superCall.append("super().__init__("); } else { - superCall.append("super(").append(klass.getName()).append(", self).__init__("); + superCall.append("super(").append(klass.getName()).append(", ").append(getSelfParameterName(origInfo)).append(").__init__("); } } else { - addComma = true; - superCall.append(superClass.getName()); - superCall.append(".__init__(self"); + addSelfToCall = true; + superCall.append(superClass.getName()).append(".__init__("); } - final StringBuilder newFunction = new StringBuilder("def __init__(self"); + final StringBuilder newFunction = new StringBuilder("def __init__("); - final Couple> couple = buildNewFunctionParamsAndSuperInitCallArgs(problemFunction, superInit); - - final List newParameters = couple.getFirst(); - if (!newParameters.isEmpty()) { - newFunction.append(", "); - } - StringUtil.join(newParameters, ", ", newFunction); + final Couple> couple = buildNewFunctionParamsAndSuperInitCallArgs(origInfo, superInfo, addSelfToCall); + StringUtil.join(couple.getFirst(), ", ", newFunction); newFunction.append(")"); + if (problemFunction.getAnnotation() != null) { newFunction.append(problemFunction.getAnnotation().getText()); } newFunction.append(":\n\t"); - final List superCallArguments = couple.getSecond(); - if (addComma && !superCallArguments.isEmpty()) { - superCall.append(", "); - } - StringUtil.join(superCallArguments, ", ", superCall); + StringUtil.join(couple.getSecond(), ", ", superCall); superCall.append(")"); final PyStatementList statementList = problemFunction.getStatementList(); @@ -130,13 +125,32 @@ public class AddCallSuperQuickFix implements LocalQuickFix { } @NotNull - private static Couple> buildNewFunctionParamsAndSuperInitCallArgs(@NotNull PyFunction origInit, - @NotNull PyFunction superInit) { + private static String getSelfParameterName(@NotNull ParametersInfo info) { + final PyParameter selfParameter = info.getSelfParameter(); + if (selfParameter == null) { + return PyNames.CANONICAL_SELF; + } + return StringUtil.defaultIfEmpty(selfParameter.getName(), PyNames.CANONICAL_SELF); + } + + @NotNull + private static Couple> buildNewFunctionParamsAndSuperInitCallArgs(@NotNull ParametersInfo origInfo, + @NotNull ParametersInfo superInfo, + boolean addSelfToCall) { final List newFunctionParams = new ArrayList(); final List superCallArgs = new ArrayList(); - final ParametersInfo origInfo = new ParametersInfo(origInit.getParameterList()); - final ParametersInfo superInfo = new ParametersInfo(superInit.getParameterList()); + final PyParameter selfParameter = origInfo.getSelfParameter(); + if (selfParameter != null && StringUtil.isNotEmpty(selfParameter.getName())) { + newFunctionParams.add(selfParameter.getText()); + } + else { + newFunctionParams.add(PyNames.CANONICAL_SELF); + } + + if (addSelfToCall) { + superCallArgs.add(getSelfParameterName(origInfo)); + } // Required parameters (not-keyword) for (PyParameter param : origInfo.getRequiredParameters()) { diff --git a/python/testData/inspections/AddCallSuperSelfNameAndAnnotationPreserved.py b/python/testData/inspections/AddCallSuperSelfNameAndAnnotationPreserved.py new file mode 100644 index 000000000000..63f31c92c99e --- /dev/null +++ b/python/testData/inspections/AddCallSuperSelfNameAndAnnotationPreserved.py @@ -0,0 +1,8 @@ +class A: + def __init__(self, x): + self.x = x + + +class B(A): + def __init__(this:'B', y): + this.y = y \ No newline at end of file diff --git a/python/testData/inspections/AddCallSuperSelfNameAndAnnotationPreserved_after.py b/python/testData/inspections/AddCallSuperSelfNameAndAnnotationPreserved_after.py new file mode 100644 index 000000000000..d84965efc1be --- /dev/null +++ b/python/testData/inspections/AddCallSuperSelfNameAndAnnotationPreserved_after.py @@ -0,0 +1,9 @@ +class A: + def __init__(self, x): + self.x = x + + +class B(A): + def __init__(this:'B', y, x): + super().__init__(x) + this.y = y \ No newline at end of file diff --git a/python/testData/inspections/AddCallSuperSelfNamePreserved.py b/python/testData/inspections/AddCallSuperSelfNamePreserved.py new file mode 100644 index 000000000000..d0064c3036ea --- /dev/null +++ b/python/testData/inspections/AddCallSuperSelfNamePreserved.py @@ -0,0 +1,8 @@ +class A: + def __init__(self, x): + self.x = x + + +class B(A): + def __init__(this, y): + this.y = y \ No newline at end of file diff --git a/python/testData/inspections/AddCallSuperSelfNamePreserved_after.py b/python/testData/inspections/AddCallSuperSelfNamePreserved_after.py new file mode 100644 index 000000000000..467395c5362f --- /dev/null +++ b/python/testData/inspections/AddCallSuperSelfNamePreserved_after.py @@ -0,0 +1,9 @@ +class A: + def __init__(self, x): + self.x = x + + +class B(A): + def __init__(this, y, x): + A.__init__(this, x) + this.y = y \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/Py3QuickFixTest.java b/python/testSrc/com/jetbrains/python/Py3QuickFixTest.java index 10ed5fa597b4..be13295c0061 100644 --- a/python/testSrc/com/jetbrains/python/Py3QuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/Py3QuickFixTest.java @@ -110,6 +110,11 @@ public class Py3QuickFixTest extends PyTestCase { }); } + // PY-16036, PY-11561 + public void testAddCallSuperSelfNameAndAnnotationPreserved() { + doInspectionTest(PyMissingConstructorInspection.class, PyBundle.message("QFIX.add.super"), true, true); + } + // PY-8991 public void testRemoveUnicodePrefixFromGluedStringNodesWithSlash() { runWithLanguageLevel(LanguageLevel.PYTHON32, new Runnable() { diff --git a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java index 1fc8572ae262..842c7fd1a940 100644 --- a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java @@ -369,6 +369,11 @@ public class PyQuickFixTest extends PyTestCase { doInspectionTest(PyMissingConstructorInspection.class, PyBundle.message("QFIX.add.super"), true, true); } + // PY-16036 + public void testAddCallSuperSelfNamePreserved() { + 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);