Infer type from issubclass

pkch · pkch · commit fe27fdaf86ba · 2017-03-14T22:07:20.000-07:00
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -5,7 +5,7 @@
 from contextlib import contextmanager
 
 from typing import (
-    Dict, Set, List, cast, Tuple, TypeVar, Union, Optional, NamedTuple, Iterator
+    Dict, Set, List, cast, Tuple, TypeVar, Union, Optional, NamedTuple, Iterator, MutableMapping
 )
 
 from mypy.errors import Errors, report_internal_error
@@ -2588,6 +2588,20 @@ def or_conditional_maps(m1: TypeMap, m2: TypeMap) -> TypeMap:
     return result
 
 
+def convert_to_types(m: MutableMapping[Expression, Type]) -> None:
+    for k in m:
+        x = m[k]
+        if isinstance(x, UnionType):
+            m[k] = UnionType([TypeType(t) for t in x.items])
+        elif isinstance(x, Instance):
+            m[k] = TypeType(m[k])
+        else:
+            # TODO: verify this is ok
+            # unknown object, don't know how to convert
+            m.clear()
+            return
+
+
 def find_isinstance_check(node: Expression,
                           type_map: Dict[Expression, Type],
                           ) -> Tuple[TypeMap, TypeMap]:
@@ -2613,6 +2627,22 @@ def find_isinstance_check(node: Expression,
                 vartype = type_map[expr]
                 type = get_isinstance_type(node.args[1], type_map)
                 return conditional_type_map(expr, vartype, type)
+        elif refers_to_fullname(node.callee, 'builtins.issubclass'):
+            expr = node.args[0]
+            if expr.literal == LITERAL_TYPE:
+                vartype = type_map[expr]
+                type = get_issubclass_type(node.args[1], type_map)
+                if isinstance(vartype, UnionType):
+                    vartype = UnionType([t.item for t in vartype.items])  # type: ignore
+                elif isinstance(vartype, TypeType):
+                    vartype = vartype.item
+                else:
+                    # TODO: verify this is ok
+                    return {}, {}  # unknown type
+                yes_map, no_map = conditional_type_map(expr, vartype, type)
+                convert_to_types(yes_map)
+                convert_to_types(no_map)
+                return yes_map, no_map
         elif refers_to_fullname(node.callee, 'builtins.callable'):
             expr = node.args[0]
             if expr.literal == LITERAL_TYPE:
@@ -2716,6 +2746,30 @@ def get_isinstance_type(expr: Expression, type_map: Dict[Expression, Type]) -> T
         return UnionType(types)
 
 
+def get_issubclass_type(expr: Expression, type_map: Dict[Expression, Type]) -> Type:
+    all_types = [type_map[e] for e in flatten(expr)]
+
+    types = []  # type: List[Type]
+
+    for type in all_types:
+        if isinstance(type, FunctionLike):
+            if type.is_type_obj():
+                # Type variables may be present -- erase them, which is the best
+                # we can do (outside disallowing them here).
+                type = erase_typevars(type.items()[0].ret_type)
+            types.append(type)
+        elif isinstance(type, TypeType):
+            types.append(type.item)
+        else:  # we didn't see an actual type, but rather a variable whose value is unknown to us
+            return None
+
+    assert len(types) != 0
+    if len(types) == 1:
+        return types[0]
+    else:
+        return UnionType(types)
+
+
 def expand_func(defn: FuncItem, map: Dict[TypeVarId, Type]) -> FuncItem:
     visitor = TypeTransformVisitor(map)
     ret = defn.accept(visitor)
