python inheritors search (doesn't work yet: Max will investigate)

python/pluginSrc/META-INF/plugin.xml

@@ -1,4 +1,7 @@
 <idea-plugin version="2" xmlns:xi="http://www.w3.org/2001/XInclude">
+  <!-- Components and extensions declared in this file work ONLY in Python plugin, not in standalone PyCharm.
+       Anything related to Python/Java integration must go in this file. -->
+
   <id>Pythonid</id>
   <name>Pythonid</name>
   <xi:include href="/META-INF/python-plugin-common.xml" xpointer="xpointer(/idea-plugin/*)"/>
@@ -6,7 +9,6 @@
   <extensions defaultExtensionNs="Pythonid">
     <importResolver implementation="com.jetbrains.python.psi.impl.PyJavaImportResolver"/>
     <typeProvider implementation="com.jetbrains.python.psi.impl.PyJavaTypeProvider"/>
-    <pySuperMethodsSearch implementation="com.jetbrains.python.psi.search.PySuperMethodsSearchExecutor"/>
     <pySuperMethodsSearch implementation="com.jetbrains.python.psi.impl.PyJavaSuperMethodsSearchExecutor"/>
   </extensions>
 </idea-plugin>

python/src/META-INF/python-plugin-common.xml

@@ -42,6 +42,7 @@
 
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyClassNameIndex"/>
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyFunctionNameIndex"/>
+    <stubIndex implementation="com.jetbrains.python.psi.stubs.PySuperClassIndex"/>
 
     <!-- Run/Debug -->
     <configurationType implementation="com.jetbrains.python.run.PythonConfigurationType"/>
@@ -60,6 +61,12 @@
   <extensionPoints>
     <extensionPoint name="importResolver" interface="com.jetbrains.python.psi.impl.PyImportResolver"/>
     <extensionPoint name="typeProvider" interface="com.jetbrains.python.psi.impl.PyTypeProvider"/>
-    <extensionPoint name="pySuperMethodsSearch" interface="com.intellij.util.QueryExecutor"/>
+    <extensionPoint qualifiedName="Pythonid.pySuperMethodsSearch" interface="com.intellij.util.QueryExecutor"/>
+    <extensionPoint qualifiedName="Pythonid.pyClassInheritorsSearch" interface="com.intellij.util.QueryExecutor"/>
   </extensionPoints>
+
+  <extensions defaultExtensionNs="Pythonid">
+    <pySuperMethodsSearch implementation="com.jetbrains.python.psi.search.PySuperMethodsSearchExecutor"/>
+    <pyClassInheritorsSearch implementation="com.jetbrains.python.psi.search.PyClassInheritorsSearchExecutor"/>
+  </extensions>
 </idea-plugin>

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

@@ -34,7 +34,7 @@ public interface PyClass extends PsiNamedElement, PyElement, PyDocStringOwner, S
   @NotNull
   PyStatementList getStatementList();
 
-  @Nullable
+  @NotNull
   PyExpression[] getSuperClassExpressions();
 
   @Nullable

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

@@ -91,7 +91,7 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     return childToPsiNotNull(PyElementTypes.STATEMENT_LIST);
   }
 
-  @Nullable
+  @NotNull
   public PyExpression[] getSuperClassExpressions() {
     final PyParenthesizedExpression superExpression = PsiTreeUtil.getChildOfType(this, PyParenthesizedExpression.class);
     if (superExpression != null) {
@@ -103,25 +103,22 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
         return new PyExpression[] { expr };
       }
     }
-    return null;
+    return PyExpression.EMPTY_ARRAY;
   }
 
   public PsiElement[] getSuperClassElements() {
     final PyExpression[] superExpressions = getSuperClassExpressions();
-    if (superExpressions != null) {
-      List<PsiElement> superClasses = new ArrayList<PsiElement>();
-      for(PyExpression expr: superExpressions) {
-        if (expr instanceof PyReferenceExpression && !PyNames.OBJECT.equals(expr.getText())) {
-          PyReferenceExpression ref = (PyReferenceExpression) expr;
-          final PsiElement result = ref.resolve();
-          if (result != null) {
-            superClasses.add(result);
-          }
+    List<PsiElement> superClasses = new ArrayList<PsiElement>();
+    for(PyExpression expr: superExpressions) {
+      if (expr instanceof PyReferenceExpression && !PyNames.OBJECT.equals(expr.getText())) {
+        PyReferenceExpression ref = (PyReferenceExpression) expr;
+        final PsiElement result = ref.resolve();
+        if (result != null) {
+          superClasses.add(result);
         }
       }
-      return superClasses.toArray(new PsiElement[superClasses.size()]);
     }
-    return null;
+    return superClasses.toArray(new PsiElement[superClasses.size()]);
   }
 
   public PyClass[] getSuperClasses() {

python/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java

@@ -10,10 +10,12 @@ import com.intellij.psi.stubs.StubElement;
 import com.intellij.util.io.DataInputOutputUtil;
 import com.intellij.util.io.PersistentStringEnumerator;
 import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.PyStubElementType;
 import com.jetbrains.python.psi.impl.PyClassImpl;
 import com.jetbrains.python.psi.stubs.PyClassNameIndex;
 import com.jetbrains.python.psi.stubs.PyClassStub;
+import com.jetbrains.python.psi.stubs.PySuperClassIndex;
 
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
@@ -33,18 +35,33 @@ public class PyClassElementType extends PyStubElementType<PyClassStub, PyClass>
   }
 
   public PyClassStub createStub(final PyClass psi, final StubElement parentStub) {
-    return new PyClassStubImpl(psi.getName(), parentStub);
+    final PyExpression[] exprs = psi.getSuperClassExpressions();
+    String[] superClasses = new String[exprs.length];
+    for(int i=0; i<exprs.length; i++) {
+      superClasses [i] = exprs [i].getText();
+    }
+    return new PyClassStubImpl(psi.getName(), parentStub, superClasses);
   }
 
   public void serialize(final PyClassStub pyClassStub, final DataOutputStream dataStream,
                         final PersistentStringEnumerator nameStorage) throws IOException {
     DataInputOutputUtil.writeNAME(dataStream, pyClassStub.getName(), nameStorage);
+    final String[] classes = pyClassStub.getSuperClasses();
+    dataStream.writeByte(classes.length);
+    for(String s: classes) {
+      DataInputOutputUtil.writeNAME(dataStream, s, nameStorage);
+    }
   }
 
   public PyClassStub deserialize(final DataInputStream dataStream, final StubElement parentStub,
                                  final PersistentStringEnumerator nameStorage) throws IOException {
     String name = DataInputOutputUtil.readNAME(dataStream, nameStorage);
-    return new PyClassStubImpl(name, parentStub);
+    int superClassCount = dataStream.readByte();
+    String[] superClasses = new String[superClassCount];
+    for(int i=0; i<superClassCount; i++) {
+      superClasses [i] = DataInputOutputUtil.readNAME(dataStream, nameStorage);
+    }
+    return new PyClassStubImpl(name, parentStub, superClasses);
   }
 
   public void indexStub(final PyClassStub stub, final IndexSink sink) {
@@ -52,5 +69,8 @@ public class PyClassElementType extends PyStubElementType<PyClassStub, PyClass>
     if (name != null) {
       sink.occurence(PyClassNameIndex.KEY, name);
     }
+    for(String s: stub.getSuperClasses()) {
+      sink.occurence(PySuperClassIndex.KEY, s);
+    }
   }
 }

python/src/com/jetbrains/python/psi/impl/stubs/PyClassStubImpl.java

@@ -11,13 +11,19 @@ import com.jetbrains.python.psi.stubs.PyClassStub;
 
 public class PyClassStubImpl extends StubBase<PyClass> implements PyClassStub {
   private final String myName;
+  private final String[] mySuperClasses;
 
-  public PyClassStubImpl(final String name, StubElement parentStub) {
+  public PyClassStubImpl(final String name, StubElement parentStub, final String[] superClasses) {
     super(parentStub, PyElementTypes.CLASS_DECLARATION);
     myName = name;
+    mySuperClasses = superClasses;
   }
 
   public String getName() {
     return myName;
   }
+
+  public String[] getSuperClasses() {
+    return mySuperClasses;
+  }
 }

python/src/com/jetbrains/python/psi/search/PyClassInheritorsSearch.java

@@ -0,0 +1,33 @@
+package com.jetbrains.python.psi.search;
+
+import com.intellij.psi.search.searches.ExtensibleQueryFactory;
+import com.intellij.util.Query;
+import com.jetbrains.python.psi.PyClass;
+
+/**
+ * @author yole
+ */
+public class PyClassInheritorsSearch extends ExtensibleQueryFactory<PyClass, PyClassInheritorsSearch.SearchParameters> {
+  public static PyClassInheritorsSearch INSTANCE = new PyClassInheritorsSearch();
+
+  public static class SearchParameters {
+    private PyClass mySuperClass;
+
+    public SearchParameters(final PyClass superClass) {
+      mySuperClass = superClass;
+    }
+
+    public PyClass getSuperClass() {
+      return mySuperClass;
+    }
+  }
+
+  private PyClassInheritorsSearch() {
+    super("Pythonid");
+  }
+
+  public static Query<PyClass> search(final PyClass superClass) {
+    final SearchParameters parameters = new SearchParameters(superClass);
+    return INSTANCE.createUniqueResultsQuery(parameters);
+  }
+}

python/src/com/jetbrains/python/psi/search/PyClassInheritorsSearchExecutor.java

@@ -0,0 +1,34 @@
+package com.jetbrains.python.psi.search;
+
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.search.ProjectScope;
+import com.intellij.psi.stubs.StubIndex;
+import com.intellij.util.Processor;
+import com.intellij.util.QueryExecutor;
+import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.stubs.PySuperClassIndex;
+
+import java.util.Collection;
+
+/**
+ * @author yole
+ */
+public class PyClassInheritorsSearchExecutor implements QueryExecutor<PyClass, PyClassInheritorsSearch.SearchParameters> {
+  public boolean execute(final PyClassInheritorsSearch.SearchParameters queryParameters, final Processor<PyClass> consumer) {
+    PyClass superClass = queryParameters.getSuperClass();
+    Project project = superClass.getProject();
+    final Collection<PyClass> candidates = StubIndex.getInstance().get(PySuperClassIndex.KEY, superClass.getName(), project,
+                                                                       ProjectScope.getAllScope(project));
+    for(PyClass candidate: candidates) {
+      final PyClass[] classes = candidate.getSuperClasses();
+      if (classes != null) {
+        for(PyClass superClassCandidate: classes) {
+          if (superClassCandidate.isEquivalentTo(superClass)) {
+            if (!consumer.process(superClassCandidate)) return false;                    
+          }
+        }
+      }
+    }
+    return true;
+  }
+}

python/src/com/jetbrains/python/psi/stubs/PyClassStub.java

@@ -7,4 +7,5 @@ import com.intellij.psi.stubs.NamedStub;
 import com.jetbrains.python.psi.PyClass;
 
 public interface PyClassStub extends NamedStub<PyClass> {
+  String[] getSuperClasses();
 }

python/src/com/jetbrains/python/psi/stubs/PySuperClassIndex.java

@@ -0,0 +1,16 @@
+package com.jetbrains.python.psi.stubs;
+
+import com.intellij.psi.stubs.StringStubIndexExtension;
+import com.intellij.psi.stubs.StubIndexKey;
+import com.jetbrains.python.psi.PyClass;
+
+/**
+ * @author yole
+ */
+public class PySuperClassIndex extends StringStubIndexExtension<PyClass> {
+  public static final StubIndexKey<String, PyClass> KEY = new StubIndexKey<String, PyClass>("Py.class.super");
+
+  public StubIndexKey<String, PyClass> getKey() {
+    return KEY;
+  }
+}

python/testData/inheritors/simple/A.py

@@ -0,0 +1,9 @@
+class A(object):
+  pass
+
+class B(A):
+  pass
+
+class C(A):
+  pass
+  

python/testSrc/com/jetbrains/python/PyInheritorsSearchTest.java

@@ -0,0 +1,41 @@
+package com.jetbrains.python;
+
+import com.intellij.codeInsight.CodeInsightTestCase;
+import com.intellij.openapi.application.PathManager;
+import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiDocumentManager;
+import com.intellij.psi.search.ProjectScope;
+import com.intellij.psi.stubs.StubIndex;
+import com.intellij.testFramework.PsiTestUtil;
+import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.search.PyClassInheritorsSearch;
+import com.jetbrains.python.psi.stubs.PyClassNameIndex;
+
+import java.util.Collection;
+
+/**
+ * @author yole
+ */
+public class PyInheritorsSearchTest extends CodeInsightTestCase {
+  public void testSimple() throws Exception {
+    setupProject();
+    final PyClass pyClass = findClass("A");
+    Collection<PyClass> inheritors = PyClassInheritorsSearch.search(pyClass).findAll();
+    assertEquals(2, inheritors.size());
+  }
+
+  private void setupProject() throws Exception {
+    String testName = getTestName(true);
+    String root = PathManager.getHomePath() + "/plugins/python/testData/inheritors/" + testName;
+    VirtualFile rootDir = PsiTestUtil.createTestProjectStructure(myProject, myModule, root, myFilesToDelete, false);
+    PsiTestUtil.addSourceContentToRoots(myModule, rootDir);
+    PsiDocumentManager.getInstance(myProject).commitAllDocuments();
+  }
+
+  private PyClass findClass(final String name) {
+    final Collection<PyClass> classes = StubIndex.getInstance().get(PyClassNameIndex.KEY, name, myProject,
+                                                                    ProjectScope.getAllScope(myProject));
+    assert classes.size() == 1;
+    return classes.iterator().next();
+  }
+}