Removed PyClassRef and PyClass.iterateAncestors() in favor of PyClass.getAncestorTypes()

python/psi-api/src/com/jetbrains/python/psi/PyClass.java

@@ -110,12 +110,6 @@ public interface PyClass extends PsiNameIdentifierOwner, PyStatement, NameDefine
    */
   boolean isNewStyleClass();
 
-  /**
-   * A lazy way to list ancestor classes width first, *not* in method-resolution order.
-   * @return an iterable of ancestor classes.
-   */
-  Iterable<PyClassRef> iterateAncestors();
-
   Iterable<PyClass> iterateAncestorClasses();
 
   /**

python/psi-api/src/com/jetbrains/python/psi/PyClassRef.java

@@ -1,99 +0,0 @@
-package com.jetbrains.python.psi;
-
-import com.intellij.psi.PsiElement;
-import com.jetbrains.python.psi.impl.PyQualifiedName;
-import com.jetbrains.python.psi.types.PyClassType;
-import org.jetbrains.annotations.Nullable;
-
-/**
- * @author yole
- */
-public class PyClassRef {
-  @Nullable private final PsiElement myElement;
-  @Nullable private final String myQName;
-  @Nullable private final PyClassType myType;
-
-  public PyClassRef(@Nullable PsiElement element) {
-    myElement = element;
-    myQName = null;
-    myType = null;
-  }
-
-  public PyClassRef(@Nullable String qName) {
-    myElement = null;
-    myQName = qName;
-    myType = null;
-  }
-
-  public PyClassRef(@Nullable PyClassType type) {
-    myElement = null;
-    myQName = null;
-    myType = type;
-  }
-
-  @Nullable
-  public PyClass getPyClass() {
-    if (myElement instanceof PyClass) {
-      return (PyClass) myElement;
-    }
-    return null;
-  }
-
-  @Nullable
-  public PsiElement getElement() {
-    return myElement;
-  }
-
-  @Nullable
-  public PyClassType getType() {
-    return myType;
-  }
-
-  @Nullable
-  public String getClassName() {
-    if (myElement instanceof PyClass) {
-      return ((PyClass)myElement).getName();
-    }
-    else if (myQName != null) {
-      final PyQualifiedName qname = PyQualifiedName.fromDottedString(myQName);
-      if (qname != null) {
-        return qname.getLastComponent();
-      }
-    }
-    else if (myType != null) {
-      return myType.getName();
-    }
-    return null;
-  }
-
-  @Nullable
-  public String getQualifiedName() {
-    if (myElement instanceof PyClass) {
-      return ((PyClass)myElement).getQualifiedName();
-    }
-    else if (myQName != null) {
-      return myQName;
-    }
-    else if (myType != null) {
-      return myType.getName();
-    }
-    return null;
-  }
-
-  @Override
-  public boolean equals(Object o) {
-    if (this == o) return true;
-    if (o == null || getClass() != o.getClass()) return false;
-
-    PyClassRef that = (PyClassRef)o;
-
-    if (myElement != null ? !myElement.equals(that.myElement) : that.myElement != null) return false;
-
-    return true;
-  }
-
-  @Override
-  public int hashCode() {
-    return myElement != null ? myElement.hashCode() : 0;
-  }
-}

python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java

@@ -27,10 +27,7 @@ import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyQualifiedName;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
-import com.jetbrains.python.psi.types.PyClassType;
-import com.jetbrains.python.psi.types.PyType;
-import com.jetbrains.python.psi.types.PyTypeParser;
-import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.psi.types.*;
 import com.jetbrains.python.toolbox.ChainIterable;
 import com.jetbrains.python.toolbox.FP;
 import org.apache.commons.httpclient.HttpClient;
@@ -316,9 +313,9 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i
       PyClass cls = inferContainingClassOf(context);
       if (cls != null) {
         String desired_name = link.substring(LINK_TYPE_PARENT.length());
-        for (PyClassRef parent : cls.iterateAncestors()) {
-          final String parent_name = parent.getClassName();
-          if (parent_name != null && parent_name.equals(desired_name)) return parent.getPyClass();
+        for (PyClass parent : cls.iterateAncestorClasses()) {
+          final String parent_name = parent.getName();
+          if (parent_name != null && parent_name.equals(desired_name)) return parent;
         }
       }
     }

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

@@ -6,6 +6,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiElementVisitor;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyClassLikeType;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -47,8 +48,11 @@ public class PyExceptionInheritInspection extends PyInspection {
           PsiElement psiElement = ((PyReferenceExpression)callee).getReference(resolveWithoutImplicits()).resolve();
           if (psiElement instanceof PyClass) {
             PyClass aClass = (PyClass) psiElement;
-            for (PyClassRef pyClass : aClass.iterateAncestors()) {
-              final String name = pyClass.getClassName();
+            for (PyClassLikeType type : aClass.getAncestorTypes(myTypeEvalContext)) {
+              if (type == null) {
+                return;
+              }
+              final String name = type.getName();
               if (name == null || "BaseException".equals(name) || "Exception".equals(name)) {
                 return;
               }

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

@@ -134,8 +134,7 @@ public class PyUnresolvedReferencesInspection extends PyInspection {
             final List<String> slots = pyClass.getSlots();
             final String attrName = node.getReferencedName();
             if (slots != null && !slots.contains(attrName) && !slots.contains(PyNames.DICT)) {
-              for (PyClassRef ref : pyClass.iterateAncestors()) {
-                final PyClass ancestor = ref.getPyClass();
+              for (PyClass ancestor : pyClass.iterateAncestorClasses()) {
                 if (ancestor == null) {
                   return;
                 }

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

@@ -618,8 +618,8 @@ public class PyUtil {
   }
 
   public static boolean hasUnresolvedAncestors(@NotNull PyClass cls) {
-    for (PyClassRef classRef : cls.iterateAncestors()) {
-      if (classRef.getPyClass() == null && classRef.getType() == null) {
+    for (PyClassLikeType type : cls.getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
+      if (type == null) {
         return true;
       }
     }
@@ -959,8 +959,10 @@ public class PyUtil {
     if (isBaseException(pyClass.getQualifiedName())) {
       return true;
     }
-    for (PyClassRef superclass : pyClass.iterateAncestors()) {
-      if (isBaseException(superclass.getQualifiedName())) return true;
+    for (PyClassLikeType type : pyClass.getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
+      if (type != null && isBaseException(type.getClassQName())) {
+        return true;
+      }
     }
     return false;
   }

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

@@ -174,21 +174,6 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     return expression;
   }
 
-  public Iterable<PyClassRef> iterateAncestors() {
-    // TODO: Change this method to getAncestorTypes()
-    // Implementation is no longer lazy, because C3 resolve for new-style classes will not be lazy
-    final List<PyClassRef> results = new ArrayList<PyClassRef>();
-    for (PyClassLikeType type : getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
-      if (type instanceof PyClassType) {
-        results.add(new PyClassRef(((PyClassType)type).getPyClass()));
-      }
-      else {
-        results.add(new PyClassRef((PyClass)null));
-      }
-    }
-    return results;
-  }
-
   @Override
   public Iterable<PyClass> iterateAncestorClasses() {
     final List<PyClass> results = new ArrayList<PyClass>();
@@ -215,8 +200,10 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     if (superClassQName.equals(getQualifiedName())) {
       return true;
     }
-    for (PyClassRef superclass : iterateAncestors()) {
-      if (superClassQName.equals(superclass.getQualifiedName())) return true;
+    for (PyClassLikeType type : getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
+      if (type != null && superClassQName.equals(type.getClassQName())) {
+        return true;
+      }
     }
     return false;
   }
@@ -274,85 +261,19 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     return PyFileImpl.getStringListFromTargetExpression(PyNames.SLOTS, getClassAttributes());
   }
 
-  @Nullable
-  private List<PyClassRef> resolveSuperClassesFromStub() {
-    final PyClassStub stub = getStub();
-    if (stub == null) {
-      return null;
-    }
-    // stub-based resolve currently works correctly only with classes in file level
-    final PsiElement parent = stub.getParentStub().getPsi();
-    if (!(parent instanceof PyFile)) {
-      // TODO[yole] handle this case
-      return null;
-    }
-
-    List<PyClassRef> result = new ArrayList<PyClassRef>();
-    for (PyQualifiedName qualifiedName : stub.getSuperClasses()) {
-      result.add(classRefFromQName((NameDefiner)parent, qualifiedName));
-    }
-    return result;
-  }
-
-  private static PyClassRef classRefFromQName(NameDefiner parent, PyQualifiedName qualifiedName) {
-    if (qualifiedName == null) {
-      return new PyClassRef((String)null);
-    }
-    NameDefiner currentParent = parent;
-    for (String component : qualifiedName.getComponents()) {
-      PsiElement element = currentParent.getElementNamed(component);
-      element = PyUtil.turnDirIntoInit(element);
-      if (element instanceof PyImportElement) {
-        element = ((PyImportElement)element).resolve();
-      }
-      if (!(element instanceof NameDefiner)) {
-        currentParent = null;
-        break;
-      }
-      currentParent = (NameDefiner)element;
-    }
-
-    if (currentParent != null) {
-      return new PyClassRef(currentParent);
-    }
-    if (qualifiedName.getComponentCount() == 1) {
-      final PyClass builtinClass = PyBuiltinCache.getInstance(parent).getClass(qualifiedName.getComponents().get(0));
-      if (builtinClass != null) {
-        return new PyClassRef(builtinClass);
-      }
-    }
-    return new PyClassRef(qualifiedName.toString());
-  }
-
   @NotNull
   public PyClass[] getSuperClasses() {
-    final PyClassStub stub = getStub();
-    if (stub != null) {
-      final List<PyClassRef> pyClasses = resolveSuperClassesFromStub();
-      if (pyClasses == null) {
-        return EMPTY_ARRAY;
-      }
-      List<PyClass> result = new ArrayList<PyClass>();
-      for (PyClassRef clsRef : pyClasses) {
-        PyClass pyClass = clsRef.getPyClass();
-        if (pyClass != null) {
-          result.add(pyClass);
-        }
-      }
-      return result.toArray(new PyClass[result.size()]);
+    final List<PyClassLikeType> superTypes = getSuperClassTypes(TypeEvalContext.fastStubOnly(null));
+    if (superTypes.isEmpty()) {
+      return EMPTY_ARRAY;
     }
-
-    PsiElement[] superClassElements = getSuperClassElements();
-    if (superClassElements.length > 0) {
-      List<PyClass> result = new ArrayList<PyClass>();
-      for (PsiElement element : superClassElements) {
-        if (element instanceof PyClass) {
-          result.add((PyClass)element);
-        }
+    final List<PyClass> result = new ArrayList<PyClass>();
+    for (PyClassLikeType type : superTypes) {
+      if (type instanceof PyClassType) {
+        result.add(((PyClassType)type).getPyClass());
       }
-      return result.toArray(new PyClass[result.size()]);
     }
-    return EMPTY_ARRAY;
+    return result.toArray(new PyClass[result.size()]);
   }
 
 
@@ -963,15 +884,17 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     final PyClass objClass = PyBuiltinCache.getInstance(this).getClass("object");
     if (this == objClass) return true; // a rare but possible case
     if (hasNewStyleMetaClass(this)) return true;
-    for (PyClassRef ancestor : iterateAncestors()) {
-      final PyClass pyClass = ancestor.getPyClass();
-      if (pyClass == null) {
+    for (PyClassLikeType type : getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
+      if (type == null) {
         // unknown, assume new-style class
         return true;
       }
-      if (pyClass == objClass) return true;
-      if (hasNewStyleMetaClass(pyClass)) {
-        return true;
+      if (type instanceof PyClassType) {
+        final PyClass pyClass = ((PyClassType)type).getPyClass();
+        if (pyClass == objClass) return true;
+        if (hasNewStyleMetaClass(pyClass)) {
+          return true;
+        }
       }
     }
     return false;

python/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java

@@ -122,6 +122,7 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
                                                              @Nullable PyExpression location,
                                                              @NotNull AccessDirection direction,
                                                              @NotNull PyResolveContext resolveContext) {
+    final TypeEvalContext context = resolveContext.getTypeEvalContext();
     PsiElement classMember = resolveByOverridingMembersProviders(this, name); //overriding members provers have priority to normal resolve
     if (classMember != null) {
       return ResolveResultList.to(classMember);
@@ -147,11 +148,11 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
       }
     }
 
-    if ("super".equals(getClassQName()) && isBuiltin(resolveContext.getTypeEvalContext()) && location instanceof PyCallExpression) {
+    if ("super".equals(getClassQName()) && isBuiltin(context) && location instanceof PyCallExpression) {
       // methods of super() call are not of class super!
       PyExpression first_arg = ((PyCallExpression)location).getArgument(0, PyExpression.class);
       if (first_arg != null) { // the usual case: first arg is the derived class that super() is proxying for
-        PyType first_arg_type = resolveContext.getTypeEvalContext().getType(first_arg);
+        PyType first_arg_type = context.getType(first_arg);
         if (first_arg_type instanceof PyClassType) {
           PyClass derived_class = ((PyClassType)first_arg_type).getPyClass();
           final Iterator<PyClass> base_it = derived_class.iterateAncestorClasses().iterator();
@@ -170,7 +171,7 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
       return ResolveResultList.to(classMember);
     }
 
-    for (PyClassLikeType type : myClass.getAncestorTypes(resolveContext.getTypeEvalContext())) {
+    for (PyClassLikeType type : myClass.getAncestorTypes(context)) {
       if (type instanceof PyClassType) {
         PsiElement superMember = resolveClassMember(((PyClassType)type).getPyClass(), myIsDefinition, name, null);
         if (superMember != null) {
@@ -201,13 +202,15 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
       return ResolveResultList.to(classMember);
     }
 
-    for (PyClassRef superClass : myClass.iterateAncestors()) {
-      final PyClass pyClass = superClass.getPyClass();
-      if (pyClass != null) {
-        PsiElement superMember = resolveByMembersProviders(new PyClassTypeImpl(pyClass, isDefinition()), name);
+    for (PyClassLikeType type : myClass.getAncestorTypes(context)) {
+      if (type instanceof PyClassType) {
+        final PyClass pyClass = ((PyClassType)type).getPyClass();
+        if (pyClass != null) {
+          PsiElement superMember = resolveByMembersProviders(new PyClassTypeImpl(pyClass, isDefinition()), name);
 
-        if (superMember != null) {
-          return ResolveResultList.to(superMember);
+          if (superMember != null) {
+            return ResolveResultList.to(superMember);
+          }
         }
       }
     }

python/src/com/jetbrains/python/testing/PythonUnitTestUtil.java

@@ -6,6 +6,8 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.util.containers.Stack;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.*;
@@ -48,11 +50,8 @@ public class PythonUnitTestUtil {
   }
 
   private static boolean isUnitTestCaseClass(PyClass cls, HashSet<String> testQualifiedNames) {
-    for (PyClassRef ancestor : cls.iterateAncestors()) {
-      if (ancestor == null) continue;
-
-      String qName = ancestor.getQualifiedName();
-      if (testQualifiedNames.contains(qName)) {
+    for (PyClassLikeType type : cls.getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
+      if (type != null && testQualifiedNames.contains(type.getClassQName())) {
         return true;
       }
     }
@@ -122,21 +121,19 @@ public class PythonUnitTestUtil {
   }
 
   public static boolean isTestCaseClass(@NotNull PyClass cls, Set<String> testQualifiedNames) {
-    for (PyClassRef ancestor : cls.iterateAncestors()) {
-      String qName = ancestor.getQualifiedName();
-      if (qName == null) continue;
-      if (testQualifiedNames.contains(qName)) {
-        return true;
-      }
-      String clsName = cls.getQualifiedName();
-      String[] names = clsName.split("\\.");
-      clsName = names[names.length - 1];
-
-      if (TEST_MATCH_PATTERN.matcher(clsName).find()) {
-        return true;
+    for (PyClassLikeType type : cls.getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
+      if (type != null) {
+        if (testQualifiedNames.contains(type.getClassQName())) {
+          return true;
+        }
+        String clsName = cls.getQualifiedName();
+        String[] names = clsName.split("\\.");
+        clsName = names[names.length - 1];
+        if (TEST_MATCH_PATTERN.matcher(clsName).find()) {
+          return true;
+        }
       }
     }
-
     return false;
   }
 }

python/src/com/jetbrains/python/testing/attest/PythonAtTestConfigurationProducer.java

@@ -5,6 +5,8 @@ import com.intellij.execution.Location;
 import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.psi.PsiElement;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.sdk.PythonSdkType;
 import com.jetbrains.python.testing.*;
 
@@ -27,9 +29,10 @@ public class PythonAtTestConfigurationProducer extends
 
   protected boolean isTestClass(PyClass pyClass) {
     if (pyClass == null) return false;
-    for (PyClassRef an : pyClass.iterateAncestors()) {
-      if ("TestBase".equals(an.getClassName()) && hasTestFunction(pyClass))
+    for (PyClassLikeType type : pyClass.getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
+      if (type != null && "TestBase".equals(type.getName()) && hasTestFunction(pyClass)) {
         return true;
+      }
     }
     return false;
   }

python/src/com/jetbrains/python/testing/pytest/PyTestUtil.java

@@ -6,6 +6,8 @@ import com.intellij.psi.PsiDirectory;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiFileSystemItem;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
 
 import java.util.HashSet;
 import java.util.List;
@@ -54,13 +56,11 @@ public class PyTestUtil {
   }
 
   public static boolean isPyTestClass(PyClass pyClass) {
-    for (PyClassRef ancestor : pyClass.iterateAncestors()) {
-      String qName = ancestor.getQualifiedName();
-      if (PYTHON_TEST_QUALIFIED_CLASSES.contains(qName)) {
+    for (PyClassLikeType type : pyClass.getAncestorTypes(TypeEvalContext.fastStubOnly(null))) {
+      if (type != null && PYTHON_TEST_QUALIFIED_CLASSES.contains(type.getClassQName())) {
         return true;
       }
     }
-
     final String className = pyClass.getName();
     if (className == null) return false;
     final String name = className.toLowerCase();

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -347,8 +347,8 @@ public class PyStubsTest extends PyTestCase {
   public void testBuiltinAncestor() {
     final PyFileImpl file = (PyFileImpl) getTestFile();
     final PyClass pyClass = file.getTopLevelClasses().get(0);
-    final PyClassRef classRef = pyClass.iterateAncestors().iterator().next();
-    assertNotNull(classRef.getPyClass());
+    final PyClass cls = pyClass.iterateAncestorClasses().iterator().next();
+    assertNotNull(cls);
     assertNotParsed(file);
   }