PyTypeParser work in progress

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

@@ -15,8 +15,8 @@ import java.util.Arrays;
 public class PyTupleType extends PyClassType implements PySubscriptableType {
   private final PyType[] myElementTypes;
 
-  public PyTupleType(PsiElement tuple, PyType[] elementTypes) {
-    super(PyBuiltinCache.getInstance(tuple).getClass("tuple"), false);
+  public PyTupleType(PsiElement anchor, PyType[] elementTypes) {
+    super(PyBuiltinCache.getInstance(anchor).getClass("tuple"), false);
     myElementTypes = elementTypes;
   }
 

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

@@ -1,9 +1,15 @@
 package com.jetbrains.python.psi.types;
 
+import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.PyFile;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.List;
+
 /**
  * @author yole
  */
@@ -16,6 +22,11 @@ public class PyTypeParser {
     if (type == null) {
       return null;
     }
+
+    if (type.startsWith("(") && type.endsWith(")")) {
+      return parseTupleType(anchor, type.substring(1, type.length()-1));
+    }
+
     final PyBuiltinCache builtinCache = PyBuiltinCache.getInstance(anchor);
 
     if (type.equals("string")) {
@@ -28,7 +39,10 @@ public class PyTypeParser {
       return builtinCache.getObjectType("dict");
     }
     if (type.startsWith("list of")) {
-      return builtinCache.getObjectType("list");
+      return parseListType(anchor, type.substring(7).trim());
+    }
+    if (type.startsWith("dict from")) {
+      return parseDictType(anchor, type.substring(9).trim());
     }
     if (type.equals("integer")) {
       return builtinCache.getIntType();
@@ -37,6 +51,42 @@ public class PyTypeParser {
     if (classType != null) {
       return classType;
     }
+
+    final PsiFile anchorFile = anchor.getContainingFile();
+    if (anchor instanceof PyFile) {
+      final PyClass aClass = ((PyFile)anchorFile).findTopLevelClass(type);
+      if (aClass != null) {
+        return new PyClassType(aClass, false);
+      }
+    }
+
     return null;
   }
+
+  private static PyType parseTupleType(PsiElement anchor, String elementTypeNames) {
+    final List<String> elements = StringUtil.split(elementTypeNames, ",");
+    PyType[] elementTypes = new PyType[elements.size()];
+    for (int i = 0; i < elementTypes.length; i++) {
+      elementTypes [i] = getTypeByName(anchor, elements.get(i).trim());
+    }
+    return new PyTupleType(anchor, elementTypes);
+  }
+
+  private static PyType parseListType(PsiElement anchor, String elementTypeName) {
+    PyClass list = PyBuiltinCache.getInstance(anchor).getClass("list");
+    PyType elementType = getTypeByName(anchor, elementTypeName);
+    return new PyCollectionTypeImpl(list, false, elementType);
+  }
+
+  @Nullable
+  private static PyType parseDictType(PsiElement anchor, String fromToTypeNames) {
+    int pos = fromToTypeNames.indexOf(" to ");
+    if (pos > 0) {
+      String toTypeName = fromToTypeNames.substring(pos + 4).trim();
+      PyClass dict = PyBuiltinCache.getInstance(anchor).getClass("dict");
+      return new PyCollectionTypeImpl(dict, false, getTypeByName(anchor, toTypeName));
+    }
+    return PyBuiltinCache.getInstance(anchor).getDictType();
+  }
+
 }

python/testData/typeParser/typeParser.py

@@ -0,0 +1,2 @@
+class MyObject:
+    pass

python/testSrc/com/jetbrains/python/PyTypeParserTest.java

@@ -0,0 +1,41 @@
+package com.jetbrains.python;
+
+import com.jetbrains.python.fixtures.PyLightFixtureTestCase;
+import com.jetbrains.python.psi.types.*;
+
+/**
+ * @author yole
+ */
+public class PyTypeParserTest extends PyLightFixtureTestCase {
+  public void testClassType() {
+    myFixture.configureByFile("typeParser/typeParser.py");
+    final PyType type = PyTypeParser.getTypeByName(myFixture.getFile(), "MyObject");
+    assertClassType(type, "MyObject");
+  }
+
+  private static void assertClassType(PyType type, final String name) {
+    assertEquals(name, ((PyClassType) type).getPyClass().getName());
+  }
+
+  public void testTupleType() {
+    myFixture.configureByFile("typeParser/typeParser.py");
+    final PyTupleType type = (PyTupleType)PyTypeParser.getTypeByName(myFixture.getFile(), "(str, MyObject)");
+    assertEquals(2, type.getElementCount());
+    assertClassType(type.getElementType(0), "str");
+    assertClassType(type.getElementType(1), "MyObject");
+  }
+
+  public void testListType() {
+    myFixture.configureByFile("typeParser/typeParser.py");
+    final PyCollectionType type = (PyCollectionType) PyTypeParser.getTypeByName(myFixture.getFile(), "list of MyObject");
+    assertClassType(type, "list");
+    assertClassType(type.getElementType(TypeEvalContext.fast()), "MyObject");
+  }
+
+  public void testDictType() {
+    myFixture.configureByFile("typeParser/typeParser.py");
+    final PyCollectionType type = (PyCollectionType) PyTypeParser.getTypeByName(myFixture.getFile(), "dict from string to MyObject");
+    assertClassType(type, "dict");
+    assertClassType(type.getElementType(TypeEvalContext.fast()), "MyObject");
+  }
+}

python/testSrc/com/jetbrains/python/PythonAllTestsSuite.java

@@ -69,7 +69,8 @@ public class PythonAllTestsSuite {
     PythonRunConfigurationTest.class,
     PyFoldingTest.class,
     EpydocStringTest.class,
-    PyEmacsTabTest.class
+    PyEmacsTabTest.class,
+    PyTypeParserTest.class
   };
 
   public static TestSuite suite() {