PY-12170 Extract Superclass: breaks code when extracting fields not defined in __init__

python/src/com/jetbrains/python/inspections/PyAssignmentToLoopOrWithParameterInspection.java

@@ -132,17 +132,7 @@ public class PyAssignmentToLoopOrWithParameterInspection extends PyInspection {
     if (element instanceof PySubscriptionExpression) {
       element = ((PySubscriptionExpression)element).getRootOperand();
     }
-    while (true) {
-      PsiReference reference = element.getReference();
-      if (reference == null) {
-        break;
-      }
-      PsiElement resolve = reference.resolve();
-      if (resolve == null || resolve.equals(element) || !PyUtil.inSameFile(resolve, element)) {
-        break;
-      }
-      element = resolve;
-    }
+    element = PyUtil.resolveToTheTop(element);
     return element;
   }
 

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

@@ -700,6 +700,38 @@ public class PyUtil {
     return PyElementGenerator.getInstance(element.getProject()).createNameIdentifier(name, LanguageLevel.forElement(element));
   }
 
+  /**
+   * Finds element declaration by resolving its references top the top but not further than file (to prevent unstubing)
+   * @param element element to resolve
+   * @return its declaration
+   */
+  @NotNull
+  public static PsiElement resolveToTheTop(@NotNull final PsiElement elementToResolve) {
+    PsiElement currentElement = elementToResolve;
+    while (true) {
+      final PsiReference reference = currentElement.getReference();
+      if (reference == null) {
+        break;
+      }
+      final PsiElement resolve = reference.resolve();
+      if ((resolve == null) || resolve.equals(currentElement) || !inSameFile(resolve, currentElement)) {
+        break;
+      }
+      currentElement = resolve;
+    }
+    return currentElement;
+  }
+
+  /**
+   * Gets class init method
+   * @param pyClass class where to find init
+   * @return class init method if any
+   */
+  @Nullable
+  public static PyFunction getInitMethod(@NotNull final PyClass pyClass) {
+    return pyClass.findMethodByName(PyNames.INIT, false);
+  }
+
   public static class KnownDecoratorProviderHolder {
     public static PyKnownDecoratorProvider[] KNOWN_DECORATOR_PROVIDERS = Extensions.getExtensions(PyKnownDecoratorProvider.EP_NAME);
 

python/src/com/jetbrains/python/refactoring/classes/membersManager/InstanceFieldsManager.java

@@ -1,5 +1,8 @@
 package com.jetbrains.python.refactoring.classes.membersManager;
 
+import com.google.common.collect.Collections2;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.NotNullPredicate;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyFunctionBuilder;
@@ -14,6 +17,9 @@ import java.util.List;
  * @author Ilya.Kazakevich
  */
 class InstanceFieldsManager extends FieldsManager {
+
+  // PY-12170
+
   InstanceFieldsManager() {
     super(false);
   }
@@ -21,37 +27,37 @@ class InstanceFieldsManager extends FieldsManager {
 
   @Override
   protected Collection<PyElement> moveAssignments(@NotNull final PyClass from,
-                                 @NotNull final Collection<PyAssignmentStatement> statements,
-                                 @NotNull final PyClass... to) {
+                                                  @NotNull final Collection<PyAssignmentStatement> statements,
+                                                  @NotNull final PyClass... to) {
     //TODO: Copy/paste with ClassFieldsManager. Move to parent?
 
     final List<PyElement> result = new ArrayList<PyElement>();
     for (final PyClass destClass : to) {
       result.addAll(copyInstanceFields(statements, destClass));
     }
-
-    deleteElements(statements);
-
-    final PyFunction fromInitMethod = from.findMethodByName(PyNames.INIT, false);
-    if (fromInitMethod != null) {
+    // Delete only declarations made in __init__ to prevent PY-12170
+    final PyFunction fromInitMethod = PyUtil.getInitMethod(from);
+    if (fromInitMethod != null) { // If class has no init method that means all its fields declared in other methods, so nothing to remove
+      deleteElements(Collections2.filter(statements, new InitsOnly(fromInitMethod)));
       //We can't leave class constructor with empty body
       PyClassRefactoringUtil.insertPassIfNeeded(fromInitMethod);
     }
     return result;
   }
 
-   /**
+  /**
    * Copies class' fields in form of assignments (instance fields) to another class.
-    * Creates init method if there is no any
+   * Creates init method if there is no any
+   *
    * @param members assignments to copy
-   * @param to destination
+   * @param to      destination
    * @return newly created fields
    */
   @NotNull
   private static List<PyAssignmentStatement> copyInstanceFields(@NotNull final Collection<PyAssignmentStatement> members,
                                                                 @NotNull final PyClass to) {
     //We need __init__ method, and if there is no any -- we need to create it
-    PyFunction toInitMethod = to.findMethodByName(PyNames.INIT, false);
+    PyFunction toInitMethod = PyUtil.getInitMethod(to);
     if (toInitMethod == null) {
       toInitMethod = createInitMethod(to);
     }
@@ -61,6 +67,7 @@ class InstanceFieldsManager extends FieldsManager {
 
   /**
    * Creates init method and adds it to certain class.
+   *
    * @param to Class where method should be added
    * @return newly created method
    */
@@ -83,4 +90,24 @@ class InstanceFieldsManager extends FieldsManager {
   protected List<PyTargetExpression> getFieldsByClass(@NotNull final PyClass pyClass) {
     return pyClass.getInstanceAttributes();
   }
+
+  private static class InitsOnly extends NotNullPredicate<PyAssignmentStatement> {
+    @NotNull
+    private final PyFunction myInitMethod;
+
+    private InitsOnly(@NotNull final PyFunction initMethod) {
+      myInitMethod = initMethod;
+    }
+
+    @Override
+    protected boolean applyNotNull(@NotNull final PyAssignmentStatement input) {
+      final PyExpression expression = input.getLeftHandSideExpression();
+      if (expression == null) {
+        return false;
+      }
+
+      final PyFunction functionWhereDeclared = PsiTreeUtil.getParentOfType(PyUtil.resolveToTheTop(expression), PyFunction.class);
+      return myInitMethod.equals(functionWhereDeclared);
+    }
+  }
 }

python/testData/refactoring/extractsuperclass/instanceNotDeclaredInInit.after.py

@@ -0,0 +1,8 @@
+class Parent(object):
+    def __init__(self):
+        self.eggs = 12
+
+
+class Child(Parent):
+    def foo(self):
+        self.eggs = 12

python/testData/refactoring/extractsuperclass/instanceNotDeclaredInInit.before.py

@@ -0,0 +1,3 @@
+class Child(object):
+    def foo(self):
+        self.eggs = 12

python/testData/refactoring/pullup/instanceNotDeclaredInInit.after.py

@@ -0,0 +1,8 @@
+class Parent(object):
+    def __init__(self):
+        self.foo = 12
+
+
+class Child(Parent):
+    def foo(self):
+        self.foo = 12

python/testData/refactoring/pullup/instanceNotDeclaredInInit.py

@@ -0,0 +1,7 @@
+class Parent(object):
+    pass
+
+
+class Child(Parent):
+    def foo(self):
+        self.foo = 12

python/testSrc/com/jetbrains/python/refactoring/classes/extractSuperclass/PyExtractSuperclassTest.java

@@ -93,6 +93,10 @@ public class PyExtractSuperclassTest extends PyClassRefactoringTest {
     doSimpleTest("Foo", "Suppa", null, true, ".foo");
   }
 
+  public void testInstanceNotDeclaredInInit() throws Exception {
+    doSimpleTest("Child", "Parent", null, true, "#eggs");
+  }
+
   public void testWithSuper() throws Exception {
     doSimpleTest("Foo", "Suppa", null, true, ".foo");
   }

python/testSrc/com/jetbrains/python/refactoring/classes/pullUp/PyPullUpTest.java

@@ -64,6 +64,10 @@ public class PyPullUpTest extends PyClassRefactoringTest {
     doHelperTest("Child", "#CLASS_VAR", "Parent");
   }
 
+  public void testInstanceNotDeclaredInInit() {
+    doHelperTest("Child", "#foo", "Parent");
+  }
+
   public void testMoveClassAttributesNoPass() {
     doHelperTest("Child2", "#CLASS_VAR", "Parent2");
   }