add field after super call

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

@@ -44,6 +44,7 @@ import com.intellij.psi.*;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.ui.awt.RelativePoint;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.IncorrectOperationException;
@@ -60,7 +61,6 @@ import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.codeInsight.stdlib.PyNamedTupleType;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
-import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.QualifiedResolveResult;
 import com.jetbrains.python.psi.types.*;
@@ -1313,13 +1313,31 @@ public class PyUtil {
       final PyStatement[] statements = statementList.getStatements();
       if (toTheBeginning && statements.length > 0) {
         final PyDocStringOwner docStringOwner = PsiTreeUtil.getParentOfType(statementList, PyDocStringOwner.class);
-        final PyStatement firstStatement = statements[0];
-        if (docStringOwner != null && firstStatement instanceof PyExpressionStatement &&
-            ((PyExpressionStatement)firstStatement).getExpression() == docStringOwner.getDocStringExpression()) {
-          element = statementList.addAfter(element, firstStatement);
+        PyStatement anchor = statements[0];
+        if (docStringOwner != null && anchor instanceof PyExpressionStatement &&
+            ((PyExpressionStatement)anchor).getExpression() == docStringOwner.getDocStringExpression()) {
+          final PyStatement next = PsiTreeUtil.getNextSiblingOfType(anchor, PyStatement.class);
+          if (next == null) {
+            return statementList.addAfter(element, anchor);
+          }
+          anchor = next;
         }
-        else
-          element = statementList.addBefore(element, firstStatement);
+        while (anchor instanceof PyExpressionStatement) {
+          final PyExpression expression = ((PyExpressionStatement)anchor).getExpression();
+          if (expression instanceof PyCallExpression) {
+            final PyExpression callee = ((PyCallExpression)expression).getCallee();
+            if ((isSuperCall((PyCallExpression)expression) || (callee != null && PyNames.INIT.equals(callee.getName())))) {
+              final PyStatement next = PsiTreeUtil.getNextSiblingOfType(anchor, PyStatement.class);
+              if (next == null) {
+                return statementList.addAfter(element, anchor);
+              }
+              anchor = next;
+            }
+            else break;
+          }
+          else break;
+        }
+        element = statementList.addBefore(element, anchor);
       }
       else {
         element = statementList.add(element);

python/testData/inspections/FieldFromUnusedParameterKeyword.py

@@ -1,3 +0,0 @@
-class A:
-    def __init__(self, <weak_warning descr="Parameter 'foo' value is not used">f<caret>oo=True</weak_warning>):
-        print('hello')

python/testData/inspections/FieldFromUnusedParameterKeyword_after.py

@@ -1,4 +0,0 @@
-class A:
-    def __init__(self, foo=True):
-        print('hello')
-        self.foo = foo

python/testData/inspections/FieldFromUnusedParameter_after.py

@@ -1,4 +0,0 @@
-class A:
-    def __init__(self, foo):
-        print('hello')
-        self.foo = foo

python/testData/quickFixes/AddFieldQuickFixTest/addFieldAddConstructor_after.py

@@ -5,8 +5,8 @@ class A:
 
 class B(A):
     def __init__(self, a, b):
-        self.x = None
         A.__init__(self, a, b)
+        self.x = None
 
     def foo(self):
         return self.x

python/testData/quickFixes/AddFieldQuickFixTest/fromUnusedParameterKeyword.py

@@ -0,0 +1,3 @@
+class A:
+    def __init__(self, fo<caret>o=True):
+        print('hello')

python/testData/quickFixes/AddFieldQuickFixTest/fromUnusedParameterKeyword_after.py

@@ -0,0 +1,4 @@
+class A:
+    def __init__(self, foo=True):
+        self.foo = foo
+        print('hello')

python/testData/quickFixes/AddFieldQuickFixTest/fromUnusedParameter_after.py

@@ -0,0 +1,4 @@
+class A:
+    def __init__(self, foo):
+        self.foo = foo
+        print('hello')

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/addPass_after.py

@@ -3,8 +3,8 @@ __author__ = 'ktisha'
 
 class A:
     def __init__(self):
-        self._a = 1
         self.b = 1
+        self._a = 1
 
     def foo(self):
         pass

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/moveToInit_after.py

@@ -3,8 +3,8 @@ __author__ = 'ktisha'
 
 class A:
     def __init__(self):
-        self._a = 1
         self.b = 1
+        self._a = 1
 
     def foo(self):
         c = 1

python/testData/refactoring/introduceField/py4437.after.py

@@ -1,7 +1,7 @@
 class SomeClass():
     def __init__(self):
-        self.x = 1
         self.a = ''
+        self.x = 1
 
     def foo(self):
         self.a

python/testSrc/com/jetbrains/python/PyQuickFixTest.java

@@ -173,13 +173,6 @@ public class PyQuickFixTest extends PyTestCase {
     doInspectionTest("AddClass.py", PyUnresolvedReferencesInspection.class, "Create class 'Xyzzy'", true, true);
   }
 
-  public void testFieldFromUnusedParameter() {  // PY-1398
-    doInspectionTest("FieldFromUnusedParameter.py", PyUnusedLocalInspection.class, "Add field 'foo' to class A", true, true);
-  }
-
-  public void testFieldFromUnusedParameterKeyword() {  // PY-1602
-    doInspectionTest("FieldFromUnusedParameterKeyword.py", PyUnusedLocalInspection.class, "Add field 'foo' to class A", true, true);
-  }
 
   public void testAddFunctionToModule() {  // PY-1602
     doInspectionTest(

python/testSrc/com/jetbrains/python/quickFixes/AddFieldQuickFixTest.java

@@ -15,13 +15,16 @@
  */
 package com.jetbrains.python.quickFixes;
 
+import com.intellij.testFramework.TestDataPath;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyQuickFixTestCase;
 import com.jetbrains.python.inspections.PyUnresolvedReferencesInspection;
+import com.jetbrains.python.inspections.PyUnusedLocalInspection;
 
 /**
  * User: ktisha
  */
+@TestDataPath("$CONTENT_ROOT/../testData//quickFixes/AddFieldQuickFixTest/")
 public class AddFieldQuickFixTest extends PyQuickFixTestCase {
 
   public void testAddClassField() {
@@ -44,4 +47,12 @@ public class AddFieldQuickFixTest extends PyQuickFixTestCase {
     doQuickFixTest(PyUnresolvedReferencesInspection.class, PyBundle.message("QFIX.NAME.add.field.$0.to.class.$1", "x", "B"));
   }
 
+  public void testFromUnusedParameter() {  // PY-1398
+    doQuickFixTest(PyUnusedLocalInspection.class, "Add field 'foo' to class A");
+  }
+
+  public void testFromUnusedParameterKeyword() {  // PY-1602
+    doQuickFixTest(PyUnusedLocalInspection.class, "Add field 'foo' to class A");
+  }
+
 }