PY-16036, PY-11561 Preserve name and annotation of "self" parameter in modified constructor

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<List<String>> couple = buildNewFunctionParamsAndSuperInitCallArgs(problemFunction, superInit);
-
-    final List<String> newParameters = couple.getFirst();
-    if (!newParameters.isEmpty()) {
-      newFunction.append(", ");
-    }
-    StringUtil.join(newParameters, ", ", newFunction);
+    final Couple<List<String>> 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<String> 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<List<String>> 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<List<String>> buildNewFunctionParamsAndSuperInitCallArgs(@NotNull ParametersInfo origInfo,
+                                                                                 @NotNull ParametersInfo superInfo,
+                                                                                 boolean addSelfToCall) {
     final List<String> newFunctionParams = new ArrayList<String>();
     final List<String> superCallArgs = new ArrayList<String>();
 
-    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()) {

python/testData/inspections/AddCallSuperSelfNameAndAnnotationPreserved.py

@@ -0,0 +1,8 @@
+class A:
+    def __init__(self, x):
+        self.x = x
+
+
+class B(A):
+    def <warning descr="Call to __init__ of super class is missed">__init_<caret>_</warning>(this:'B', y):
+        this.y = y

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

python/testData/inspections/AddCallSuperSelfNamePreserved.py

@@ -0,0 +1,8 @@
+class A:
+    def __init__(self, x):
+        self.x = x
+
+
+class B(A):
+    def <warning descr="Call to __init__ of super class is missed">__init_<caret>_</warning>(this, y):
+        this.y = y

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

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() {

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);