Building new repository. Not using it yet though.

python/src/com/jetbrains/python/PythonLanguage.java

@@ -4,9 +4,19 @@ import com.intellij.lang.Language;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiFileFactory;
+import com.intellij.psi.stubs.SerializationManager;
 import com.intellij.psi.tree.IFileElementType;
-import com.jetbrains.python.psi.impl.PyElementGeneratorImpl;
+import com.intellij.psi.tree.IStubFileElementType;
 import com.jetbrains.python.psi.PyElementGenerator;
+import com.jetbrains.python.psi.impl.PyElementGeneratorImpl;
+import com.jetbrains.python.psi.impl.stubs.PyClassSerializer;
+import com.jetbrains.python.psi.impl.stubs.PyFunctionSerializer;
+import com.jetbrains.python.psi.impl.stubs.PyParameterListSerializer;
+import com.jetbrains.python.psi.impl.stubs.PyParameterSerializer;
+import com.jetbrains.python.psi.stubs.PyClassStub;
+import com.jetbrains.python.psi.stubs.PyFunctionStub;
+import com.jetbrains.python.psi.stubs.PyParameterListStub;
+import com.jetbrains.python.psi.stubs.PyParameterStub;
 import com.jetbrains.python.validation.*;
 
 import java.util.Set;
@@ -17,8 +27,8 @@ import java.util.concurrent.CopyOnWriteArraySet;
  */
 public class PythonLanguage extends Language {
   private final PyElementGenerator elementGenerator = new PyElementGeneratorImpl(this);
-  public final IFileElementType ELTYPE_FILE = new IFileElementType(this);
-  protected Set<Class<? extends PyAnnotator>> _annotators = new CopyOnWriteArraySet<Class<? extends PyAnnotator>>();
+  private final IFileElementType ELTYPE_FILE = new IStubFileElementType(this);
+  private final Set<Class<? extends PyAnnotator>> _annotators = new CopyOnWriteArraySet<Class<? extends PyAnnotator>>();
 
   public static PythonLanguage getInstance() {
     return (PythonLanguage) PythonFileType.INSTANCE.getLanguage();
@@ -37,6 +47,13 @@ public class PythonLanguage extends Language {
     _annotators.add(UnresolvedReferenceAnnotator.class);
   }
 
+  {
+    SerializationManager.getInstance().registerSerializer(PyClassStub.class, new PyClassSerializer());
+    SerializationManager.getInstance().registerSerializer(PyFunctionStub.class, new PyFunctionSerializer());
+    SerializationManager.getInstance().registerSerializer(PyParameterStub.class, new PyParameterSerializer());
+    SerializationManager.getInstance().registerSerializer(PyParameterListStub.class, new PyParameterListSerializer());
+  }
+
   protected PythonLanguage() {
     super("Python");
   }

python/src/com/jetbrains/python/PythonParserDefinition.java

@@ -40,10 +40,10 @@ import org.jetbrains.annotations.NotNull;
  * @author Keith Lea
  */
 public class PythonParserDefinition implements ParserDefinition {
-  private PythonLanguage language;
+  private final PythonLanguage language;
 
-  private TokenSet myWhitespaceTokens;
-  private TokenSet myCommentTokens;
+  private final TokenSet myWhitespaceTokens;
+  private final TokenSet myCommentTokens;
 
   public PythonParserDefinition() {
     language = (PythonLanguage)PythonFileType.INSTANCE.getLanguage();

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

@@ -20,4 +20,7 @@ public interface PyFile extends PyElement, PsiFile {
     List<PyStatement> getStatements();
 
     PythonLanguage getPyLanguage();
+
+    List<PyClass> getTopLevelClasses();
+    List<PyFunction> getTopLevelFunctions();
 }

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

@@ -23,10 +23,13 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiElementVisitor;
 import com.intellij.psi.ResolveState;
 import com.intellij.psi.scope.PsiScopeProcessor;
+import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.tree.TokenSet;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.PythonLanguage;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.stubs.PyClassStub;
+import com.jetbrains.python.psi.stubs.PyFunctionStub;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -108,4 +111,48 @@ public class PyFileImpl extends PsiFileBase implements PyFile {
   public PythonLanguage getPyLanguage() {
     return (PythonLanguage)getLanguage();
   }
+
+  public List<PyClass> getTopLevelClasses() {
+    List<PyClass> result = new ArrayList<PyClass>();
+    final StubElement stub = getStub();
+    if (stub != null) {
+      final List<StubElement> children = stub.getChildrenStubs();
+      for (StubElement child : children) {
+        if (child instanceof PyClassStub) {
+          result.add(((PyClassStub)child).getPsi());
+        }
+      }
+    }
+    else {
+      for (PsiElement child : getChildren()) {
+        if (child instanceof PyClass) {
+          result.add((PyClass)child);
+        }
+      }
+    }
+
+    return result;
+  }
+
+  public List<PyFunction> getTopLevelFunctions() {
+    List<PyFunction> result = new ArrayList<PyFunction>();
+    final StubElement stub = getStub();
+    if (stub != null) {
+      final List<StubElement> children = stub.getChildrenStubs();
+      for (StubElement child : children) {
+        if (child instanceof PyFunctionStub) {
+          result.add(((PyFunctionStub)child).getPsi());
+        }
+      }
+    }
+    else {
+      for (PsiElement child : getChildren()) {
+        if (child instanceof PyFunction) {
+          result.add((PyFunction)child);
+        }
+      }
+    }
+    
+    return result;
+  }
 }

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

@@ -0,0 +1,7 @@
+/*
+ * @author max
+ */
+package com.jetbrains.python.psi.impl.stubs;
+
+public class PyFileSerializer {
+}

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

@@ -30,7 +30,7 @@ public class PyFunctionStubImpl extends StubBase<PyFunction> implements PyFuncti
 
   @NotNull
   public PyParameterListStub getParameterList() {
-    for (StubElement childStub : getChildStubs()) {
+    for (StubElement childStub : getChildrenStubs()) {
       if (childStub.getStubType() == PyElementTypes.PARAMETER_LIST) {
         return (PyParameterListStub)childStub;
       }

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

@@ -23,7 +23,7 @@ public class PyParameterListStubImpl extends StubBase<PyParameterList> implement
   }
 
   public PyParameterStub[] getParameters() {
-    final List<StubElement> stubs = getChildStubs();
+    final List<StubElement> stubs = getChildrenStubs();
     PyParameterStub[] result = new PyParameterStub[stubs.size()];
 
     for (int i = 0; i < stubs.size(); i++) {

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

@@ -4,10 +4,7 @@
 package com.jetbrains.python.psi.stubs;
 
 import com.intellij.psi.stubs.NamedStub;
-import com.intellij.psi.stubs.SerializerClass;
 import com.jetbrains.python.psi.PyClass;
-import com.jetbrains.python.psi.impl.stubs.PyClassSerializer;
 
-@SerializerClass(PyClassSerializer.class)
 public interface PyClassStub extends NamedStub<PyClass> {
 }