Conditional type-check based on callable call (#2627)

Fixes #1973

Author: afrieder

mypy/checker.py

@@ -2507,6 +2507,76 @@ def conditional_type_map(expr: Expression,
         return {}, {}
 
 
+def partition_by_callable(type: Optional[Type]) -> Tuple[List[Type], List[Type]]:
+    """Takes in a type and partitions that type into callable subtypes and
+    uncallable subtypes.
+
+    Thus, given:
+    `callables, uncallables = partition_by_callable(type)`
+
+    If we assert `callable(type)` then `type` has type Union[*callables], and
+    If we assert `not callable(type)` then `type` has type Union[*uncallables]
+
+    Guaranteed to not return [], []"""
+    if isinstance(type, FunctionLike) or isinstance(type, TypeType):
+        return [type], []
+
+    if isinstance(type, AnyType):
+        return [type], [type]
+
+    if isinstance(type, UnionType):
+        callables = []
+        uncallables = []
+        for subtype in type.items:
+            subcallables, subuncallables = partition_by_callable(subtype)
+            callables.extend(subcallables)
+            uncallables.extend(subuncallables)
+        return callables, uncallables
+
+    if isinstance(type, TypeVarType):
+        return partition_by_callable(type.erase_to_union_or_bound())
+
+    if isinstance(type, Instance):
+        method = type.type.get_method('__call__')
+        if method:
+            callables, uncallables = partition_by_callable(method.type)
+            if len(callables) and not len(uncallables):
+                # Only consider the type callable if its __call__ method is
+                # definitely callable.
+                return [type], []
+        return [], [type]
+
+    return [], [type]
+
+
+def conditional_callable_type_map(expr: Expression,
+                                  current_type: Optional[Type],
+                                  ) -> Tuple[TypeMap, TypeMap]:
+    """Takes in an expression and the current type of the expression.
+
+    Returns a 2-tuple: The first element is a map from the expression to
+    the restricted type if it were callable. The second element is a
+    map from the expression to the type it would hold if it weren't
+    callable."""
+    if not current_type:
+        return {}, {}
+
+    if isinstance(current_type, AnyType):
+        return {}, {}
+
+    callables, uncallables = partition_by_callable(current_type)
+
+    if len(callables) and len(uncallables):
+        callable_map = {expr: UnionType.make_union(callables)} if len(callables) else None
+        uncallable_map = {expr: UnionType.make_union(uncallables)} if len(uncallables) else None
+        return callable_map, uncallable_map
+
+    elif len(callables):
+        return {}, None
+
+    return None, {}
+
+
 def is_true_literal(n: Expression) -> bool:
     return (refers_to_fullname(n, 'builtins.True')
             or isinstance(n, IntExpr) and n.value == 1)
@@ -2580,7 +2650,7 @@ def find_isinstance_check(node: Expression,
                           type_map: Dict[Expression, Type],
                           ) -> Tuple[TypeMap, TypeMap]:
     """Find any isinstance checks (within a chain of ands).  Includes
-    implicit and explicit checks for None.
+    implicit and explicit checks for None and calls to callable.
 
     Return value is a map of variables to their types if the condition
     is true and a map of variables to their types if the condition is false.
@@ -2601,6 +2671,11 @@ 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.callable'):
+            expr = node.args[0]
+            if expr.literal == LITERAL_TYPE:
+                vartype = type_map[expr]
+                return conditional_callable_type_map(expr, vartype)
     elif (isinstance(node, ComparisonExpr) and experiments.STRICT_OPTIONAL):
         # Check for `x is None` and `x is not None`.
         is_not = node.operators == ['is not']

mypy/test/testcheck.py

@@ -37,6 +37,7 @@
 ]
 fast_parser_files = [
     'check-basic.test',
+    'check-callable.test',
     'check-classes.test',
     'check-statements.test',
     'check-generics.test',