PY-21220 Handle moved classes that contain both __init__ and __new__

I changed behavior of PyClassImpl#NameFinder processor, so that it
always tries to find an element with the first of the names passed to
its constructor. In particular, #findInitOrNew() returns __init__
unless there was only __new__ defined in the class. Otherwise its
behavior contradicts with the method's javadoc.

python/src/com/jetbrains/python/psi/impl/PyClassImpl.java

@@ -527,6 +527,8 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
   private static class NameFinder<T extends PyElement> implements Processor<T> {
     private T myResult;
     private final String[] myNames;
+    private int myLastResultIndex = -1;
+    private PyClass myLastVisitedClass = null;
 
     public NameFinder(String... names) {
       myNames = names;
@@ -537,11 +539,26 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
       return myResult;
     }
 
+    @Nullable
+    protected PyClass getContainingClass(@NotNull T element) {
+      return null;
+    }
+
     public boolean process(T target) {
-      final String targetName = target.getName();
-      for (String name : myNames) {
-        if (name.equals(targetName)) {
-          myResult = target;
+      final PyClass currentClass = getContainingClass(target);
+      // Stop when the current class changes and there was a result
+      if (myLastVisitedClass != null && myLastVisitedClass != currentClass && myResult != null) {
+        return false;
+      }
+
+      myLastVisitedClass = currentClass;
+
+      final int index = ArrayUtil.indexOf(myNames, target.getName());
+      // Do not depend on the order in which elements appear, always try to find the first one
+      if (index >= 0 && (myLastResultIndex == -1 || index < myLastResultIndex)) {
+        myLastResultIndex = index;
+        myResult = target;
+        if (index == 0) {
           return false;
         }
       }
@@ -570,7 +587,13 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
   public PyFunction findInitOrNew(boolean inherited, final @Nullable TypeEvalContext context) {
     NameFinder<PyFunction> proc;
     if (isNewStyleClass(context)) {
-      proc = new NameFinder<>(PyNames.INIT, PyNames.NEW);
+      proc = new NameFinder<PyFunction>(PyNames.INIT, PyNames.NEW) {
+        @Nullable
+        @Override
+        protected PyClass getContainingClass(@NotNull PyFunction element) {
+          return element.getContainingClass();
+        }
+      };
     }
     else {
       proc = new NameFinder<>(PyNames.INIT);

python/testData/refactoring/move/referenceToClassWithNewInMovedSymbol/after/src/classFile.py

@@ -0,0 +1,10 @@
+from collections import namedtuple
+
+
+class Pipeline(namedtuple('_Pipeline', 'name')):
+
+    def __new__(cls, name):
+        return super(Pipeline, cls).__new__(cls, name)
+
+    def __init__(self, name):
+        pass

python/testData/refactoring/move/referenceToClassWithNewInMovedSymbol/after/src/fromFile.py

@@ -0,0 +1,3 @@
+
+
+

python/testData/refactoring/move/referenceToClassWithNewInMovedSymbol/after/src/toFile.py

@@ -0,0 +1,5 @@
+from classFile import Pipeline
+
+
+def fnToMove():
+    return Pipeline('name')

python/testData/refactoring/move/referenceToClassWithNewInMovedSymbol/before/src/classFile.py

@@ -0,0 +1,10 @@
+from collections import namedtuple
+
+
+class Pipeline(namedtuple('_Pipeline', 'name')):
+
+    def __new__(cls, name):
+        return super(Pipeline, cls).__new__(cls, name)
+
+    def __init__(self, name):
+        pass

python/testData/refactoring/move/referenceToClassWithNewInMovedSymbol/before/src/fromFile.py

@@ -0,0 +1,5 @@
+from classFile import Pipeline
+
+
+def fnToMove():
+    return Pipeline('name')

python/testData/resolve/InitOrNewReturnsInitWhenNewIsFirst.py

@@ -0,0 +1,6 @@
+class MyClass(object):
+    def __new__(cls):
+        return object.__new__(cls)
+
+    def __init__(self):
+        self.x = 42

python/testData/resolve/InitOrNewReturnsInitWhenNewIsLast.py

@@ -0,0 +1,6 @@
+class MyClass(object):
+    def __init__(self):
+        self.x = 42
+
+    def __new__(cls):
+        return object.__new__(cls)

python/testSrc/com/jetbrains/python/PyResolveTest.java

@@ -30,6 +30,7 @@ import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
 import com.jetbrains.python.psi.resolve.ImportedResolveResult;
+import com.jetbrains.python.psi.types.TypeEvalContext;
 
 public class PyResolveTest extends PyResolveTestCase {
   @Override
@@ -70,6 +71,22 @@ public class PyResolveTest extends PyResolveTestCase {
     assertEquals(PyNames.INIT, ((PyFunction)target).getName());
   }
 
+  public void testInitOrNewReturnsInitWhenNewIsFirst() {
+    doTestInitOrNewReturnsInit();
+  }
+  
+  public void testInitOrNewReturnsInitWhenNewIsLast() {
+    doTestInitOrNewReturnsInit();
+  }
+
+  private void doTestInitOrNewReturnsInit() {
+    myFixture.configureByFile("resolve/" + getTestName(false) + ".py");
+    final PyClass pyClass = PsiTreeUtil.findChildOfType(myFixture.getFile(), PyClass.class);
+    assertNotNull(pyClass);
+    final PyFunction init = pyClass.findInitOrNew(false, TypeEvalContext.userInitiated(myFixture.getProject(), myFixture.getFile()));
+    assertEquals(PyNames.INIT, init.getName());
+  }
+
   public void testToConstructorInherited() {
     ResolveResult[] targets = multiResolve();
     assertEquals(2, targets.length); // to class, to init

python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java

@@ -376,6 +376,11 @@ public class PyMoveTest extends PyTestCase {
     doMoveSymbolTest("FOO", "b.py");
   }
 
+  // PY-21220
+  public void testReferenceToClassWithNewInMovedSymbol() {
+    doMoveSymbolTest("fnToMove", "toFile.py");
+  }
+
   private void doMoveFileTest(String fileName, String toDirName) {
     Project project = myFixture.getProject();
     PsiManager manager = PsiManager.getInstance(project);