From 41aaebef6908b6d77c5c50534fd7cc54c6876e1a Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Mon, 13 Feb 2017 17:05:46 +0300
Subject: [PATCH] PY-20057 Initial support of typing.Type (types of direct
 classes)

---
 .../typing/PyTypingTypeProvider.java          | 24 +++++++++++++
 .../documentation/PyTypeModelBuilder.java     | 25 +++++++++++++
 .../com/jetbrains/python/PyTypingTest.java    | 36 +++++++++++++++++++
 3 files changed, 85 insertions(+)

diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
index 6ee353183b47..d1de08405a58 100644
--- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
@@ -62,6 +62,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
   public static final String COROUTINE = "typing.Coroutine";
   public static final String NAMEDTUPLE = "typing.NamedTuple";
   public static final String GENERIC = "typing.Generic";
+  public static final String TYPE = "typing.Type";
 
   public static final Pattern TYPE_COMMENT_PATTERN = Pattern.compile("# *type: *(.*)");
 
@@ -482,6 +483,10 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
       if (genericType != null) {
         return Ref.create(genericType);
       }
+      final PyType classObjType = getClassObjectType(resolved, context);
+      if (classObjType != null) {
+        return Ref.create(classObjType);
+      }
       final Ref<PyType> classType = getClassType(resolved, context.getTypeContext());
       if (classType != null) {
         return classType;
@@ -501,6 +506,25 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     }
   }
 
+  @Nullable
+  private static PyType getClassObjectType(@Nullable PsiElement resolved, @NotNull Context context) {
+    if (resolved instanceof PySubscriptionExpression) {
+      final PySubscriptionExpression subsExpr = (PySubscriptionExpression)resolved;
+      final PyExpression operand = subsExpr.getOperand();
+      final Collection<String> operandNames = resolveToQualifiedNames(operand, context.getTypeContext());
+      if (operandNames.contains(TYPE)) {
+        final PyExpression indexExpr = subsExpr.getIndexExpression();
+        if (indexExpr != null) {
+          final PyType type = Ref.deref(getType(indexExpr, context));
+          if (type instanceof PyClassType && !((PyClassType)type).isDefinition()) {
+            return new PyClassTypeImpl(((PyClassType)type).getPyClass(), false);
+          }
+        }
+      }
+    }
+    return null;
+  }
+
   @Nullable
   private static Ref<PyType> getAnyType(@NotNull PsiElement element) {
     final PyQualifiedNameOwner qualifiedNameOwner = as(element, PyQualifiedNameOwner.class);
diff --git a/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java b/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java
index 7fc8cf0a6a6b..2f39404ecfa8 100644
--- a/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java
+++ b/python/src/com/jetbrains/python/documentation/PyTypeModelBuilder.java
@@ -186,6 +186,19 @@ public class PyTypeModelBuilder {
       visitor.param(this);
     }
   }
+  
+  static class ClassObjectType extends TypeModel {
+    private final TypeModel classType;
+
+    public ClassObjectType(TypeModel classType) {
+      this.classType = classType;
+    }
+
+    @Override
+    void accept(@NotNull TypeVisitor visitor) {
+      visitor.classObject(this);
+    }
+  } 
 
   /**
    * Builds tree-like type model for PyType
@@ -251,6 +264,9 @@ public class PyTypeModelBuilder {
     else if (type instanceof PyCallableType && !(type instanceof PyClassLikeType)) {
       result = build((PyCallableType)type);
     }
+    else if (type instanceof PyClassType && !((PyClassType)type).isDefinition()) {
+      result = new ClassObjectType(build(type, allowUnions));
+    }
     if (result == null) {
       result = NamedType.nameOrAny(type);
     }
@@ -309,6 +325,8 @@ public class PyTypeModelBuilder {
     void optional(OptionalType type);
 
     void tuple(TupleType type);
+
+    void classObject(ClassObjectType type);
   }
 
   private static class TypeToStringVisitor extends TypeNameVisitor {
@@ -503,5 +521,12 @@ public class PyTypeModelBuilder {
       }
       add("]");
     }
+
+    @Override
+    public void classObject(ClassObjectType type) {
+      add("Type[");
+      type.classType.accept(this);
+      add("]");
+    }
   }
 }
diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java
index 8a7e76ad6ffc..4b64e1e9e884 100644
--- a/python/testSrc/com/jetbrains/python/PyTypingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java
@@ -783,6 +783,42 @@ public class PyTypingTest extends PyTestCase {
            "expr = C(0).get()\n");
   }
 
+  // PY-20057
+  public void testClassObjectType() {
+    doTest("Type[MyClass]",
+           "from typing import Type\n" +
+           "\n" +
+           "class MyClass:\n" +
+           "    pass\n" +
+           "\n" +
+           "def f(x: Type[MyClass]): \n" +
+           "    expr = x");
+  }
+  
+  // PY-20057
+  public void testConstrainedClassObjectTypeOfParam() {
+    doTest("Type[T]",
+           "from typing import Type, TypeVar\n" +
+           "\n" +
+           "T = TypeVar('T', bound=int)\n" +
+           "\n" +
+           "def f(x: Type[T]):\n" +
+           "    expr = x");
+  }
+  
+  // PY-20057
+  public void testConstrainedClassObjectTypeOfResultValue() {
+    doTest("int",
+           "from typing import Type, TypeVar\n" +
+           "\n" +
+           "T = TypeVar('T', bound=int)\n" +
+           "\n" +
+           "def f(x: Type[T]) -> T:\n" +
+           "    return x()\n" +
+           "\n" +
+           "expr = f(int)");
+  }
+
   private void doTestNoInjectedText(@NotNull String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);
     final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject());