Multi-resolve at a constructor call resolves to constructor, not class, or both class and constructor

if the constructor is inherited.

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

@@ -22,6 +22,7 @@ import com.intellij.psi.PsiReference;
 import com.intellij.psi.ResolveResult;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.types.PyClassType;
 import com.jetbrains.python.psi.types.PyType;
@@ -82,11 +83,9 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress
           EnumSet<Flag> flags = EnumSet.noneOf(Flag.class);
           //boolean is_inst = isByInstance();
           if (isByInstance()) flags.add(Flag.IMPLICIT_FIRST_ARG);
-          if (resolved instanceof PyClass) { // constructor call
-            final PyClass cls = (PyClass)resolved;
-            resolved = cls.findMethodByName("__init__");  // XXX move this name to PyNames
-            //is_inst |= true;
-            flags.add(Flag.IMPLICIT_FIRST_ARG);
+          if (resolved instanceof PyFunction) {
+            PyFunction meth = (PyFunction)resolved; // constructor call?
+            if (PyNames.INIT.equals(meth.getName())) flags.add(Flag.IMPLICIT_FIRST_ARG);
           }
           if (resolved != null) {
             // look for closest decorator

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

@@ -260,15 +260,56 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
     else return EMPTY_LIST; 
   }
 
+  // sorts and modifies results of resolveInner
   private ResolveResult[] multiResolveInner(boolean incomplete) {
     final String referencedName = getReferencedName();
     if (referencedName == null) return ResolveResult.EMPTY_ARRAY;
 
-    // crude logic right here to see it work
-
     List<RatedResolveResult> targets = resolveInner();
     if (targets.size() == 0) return ResolveResult.EMPTY_ARRAY;
 
+    // change class results to constructor results if there are any
+    if (getParent() instanceof PyCallExpression) { // we're a call
+      ListIterator<RatedResolveResult> it = targets.listIterator();
+      while (it.hasNext()) {
+        final RatedResolveResult rrr = it.next();
+        final PsiElement elt = rrr.getElement();
+        if (elt instanceof PyClass) {
+          PyClass cls = (PyClass)elt;
+          PyFunction init = cls.findMethodByName(PyNames.INIT);
+          if (init != null) {
+            // replace
+            final PyFunction the_init = init;
+            it.set(new RatedResolveResult(){
+              public int getRate() { return rrr.getRate(); }
+
+              public PsiElement getElement() { return the_init; }
+
+              public boolean isValidResult() { return true; }
+            });
+          }
+          else { // init not found; maybe it's ancestor's
+            for (PyClass ancestor : cls.iterateAncestors()) {
+              init = ancestor.findMethodByName(PyNames.INIT);
+              if (init != null) {
+                final PyFunction the_init = init;
+                // add to resuls as low priority
+                it.add(new RatedResolveResult(){
+                  public int getRate() { return RATE_LOW; }
+
+                  public PsiElement getElement() { return the_init; }
+
+                  public boolean isValidResult() { return true; }
+                });
+                break;
+              }
+            }
+          }
+        }
+      }
+    }
+
+    // put everything in a sorting container
     List<RatedResolveResult> ret = new SortedList<RatedResolveResult>(new Comparator<RatedResolveResult>() {
       public int compare(final RatedResolveResult one, final RatedResolveResult another) {
         return another.getRate() - one.getRate();
@@ -280,7 +321,7 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
   }
 
   /**
-   * Reso   lves reference to possible referred elements.
+   * Resolves reference to possible referred elements.
    * First element is always what resolve() would return.
    * Imported module names: to module file, or {directory, '___init__.py}' for a qualifier.
    * @todo Local identifiers: a list of definitions in the most recent compound statement 
@@ -300,6 +341,7 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
     }
   }
 
+
   public String getCanonicalText() {
     return null;
   }

python/testData/resolve/ToConstructor.py

@@ -0,0 +1,6 @@
+# must resolve to constructor
+class Foo:
+  def __init__(self):
+    pass
+
+F<ref>oo()

python/testData/resolve/ToConstructorInherited.py

@@ -0,0 +1,9 @@
+# multi-resolve: to class and to inherited constructor
+class Foo:
+  def __init__(self):
+    pass
+
+class Bar(Foo):
+  pass
+
+<ref>Bar()

python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java

@@ -18,14 +18,6 @@ public class PyMultiFileResolveTest extends CodeInsightTestCase {
     assertEquals("ImportedFile.py", ((PyFile) element).getName());
   }
 
-  /*
-  public void testFromImport() throws Exception {
-    PsiElement element = doResolve();
-    assertTrue(element instanceof PyFunction);
-    assertEquals("func", ((PyFunction) element).getName());
-  }
-  */
-
   public void testFromImport() throws Exception {
     ResolveResult[] results = doMultiResolve();
     assertTrue(results.length == 2); // func and import stmt
@@ -36,14 +28,6 @@ public class PyMultiFileResolveTest extends CodeInsightTestCase {
     assertTrue("is import?", import_elt instanceof PyImportElement);
   }
 
-  /*
-  public void testFromImportStar() throws Exception {
-    PsiElement element = doResolve();
-    assertTrue(element instanceof PyFunction);
-    assertEquals("func", ((PyFunction) element).getName());
-  }
-  */
-
   public void testFromImportStar() throws Exception {
     ResolveResult[] results = doMultiResolve();
     assertTrue(results.length == 2); // func and import-* stmt
@@ -101,13 +85,6 @@ public class PyMultiFileResolveTest extends CodeInsightTestCase {
     assertTrue(element instanceof PyFunction);
   }
 
-  /*
-  public void testTransitiveImport() throws Exception {
-    PsiElement element = doResolve();
-    assertTrue(element instanceof PyTargetExpression);
-  }
-  */
-
   public void testTransitiveImport() throws Exception {
     ResolveResult[] results = doMultiResolve();
     assertTrue(results.length == 2); // func and import stmt
@@ -115,13 +92,6 @@ public class PyMultiFileResolveTest extends CodeInsightTestCase {
     assertTrue("is target?", elt instanceof PyTargetExpression);
   }
 
-  /*
-  public void testResolveInPkg() throws Exception {
-    PsiElement element = doResolve();
-    assertTrue(element instanceof PyFunction);
-  }
-  */
-
   public void testResolveInPkg() throws Exception {
     ResolveResult[] results = doMultiResolve();
     assertTrue(results.length == 2); // func and import stmt

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

@@ -7,6 +7,8 @@ package com.jetbrains.python;
 import com.intellij.openapi.application.PathManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiReference;
+import com.intellij.psi.ResolveResult;
+import com.intellij.psi.PsiPolyVariantReference;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.testFramework.ResolveTestCase;
 import com.jetbrains.python.psi.*;
@@ -17,6 +19,12 @@ public class PyResolveTest extends ResolveTestCase {
     return ref.resolve();
   }
 
+  private ResolveResult[] multiResolve() throws Exception {
+    PsiReference ref = configureByFile(getTestName(false) + ".py");
+    assertTrue(ref instanceof PsiPolyVariantReference);
+    return ((PsiPolyVariantReference)ref).multiResolve(false);
+  }
+
   public void testClass() throws Exception {
     PsiElement target = resolve();
     assertTrue(target instanceof PyClass);
@@ -27,7 +35,31 @@ public class PyResolveTest extends ResolveTestCase {
     assertTrue(targetElement instanceof PyFunction);
   }
 
-  // NOTE: maybe this test does not belong exactly here; still it's the best place currently. 
+  public void testToConstructor() throws Exception {
+    PsiElement target = resolve();
+    assertTrue(target instanceof PyFunction);
+    assertEquals(((PyFunction)target).getName(), PyNames.INIT);
+  }
+
+  public void testToConstructorInherited() throws Exception {
+    ResolveResult[] targets = multiResolve();
+    assertEquals(targets.length, 2); // to class, to init
+    PsiElement elt;
+    // class
+    elt = targets[0].getElement();
+    assertTrue(elt instanceof PyClass);
+    assertEquals(((PyClass)elt).getName(), "Bar");
+    // init
+    elt = targets[1].getElement();
+    assertTrue(elt instanceof PyFunction);
+    PyFunction fun = (PyFunction)elt;
+    assertEquals(fun.getName(), PyNames.INIT);
+    PyClass cls = fun.getContainingClass();
+    assertNotNull(cls);
+    assertEquals(cls.getName(), "Foo");
+  }
+
+  // NOTE: maybe this test does not belong exactly here; still it's the best place currently.
   public void testComplexCallee() throws Exception {
     PsiElement targetElement = resolve();
     PyExpression assigned = ((PyAssignmentStatement) targetElement.getContext()).getAssignedValue();