Refine parent type when narrowing "lookup" expressions (#7917)

This diff adds support for the following pattern:

```python
from typing import Enum, List
from enum import Enum

class Key(Enum):
    A = 1
    B = 2

class Foo:
    key: Literal[Key.A]
    blah: List[int]

class Bar:
    key: Literal[Key.B]
    something: List[str]

x: Union[Foo, Bar]
if x.key is Key.A:
    reveal_type(x)  # Revealed type is 'Foo'
else:
    reveal_type(x)  # Revealed type is 'Bar'
```

In short, when we do `x.key is Key.A`, we "propagate" the information
we discovered about `x.key` up to refine the type of `x`.

We perform this propagation only when `x` is a Union and only when we
are doing member or index lookups into instances, typeddicts,
namedtuples, and tuples. For indexing operations, we have one additional
limitation: we *must* use a literal expression in order for narrowing
to work at all. Using Literal types or Final instances won't work;
See #7905 for more details.

To put it another way, this adds support for tagged unions, I guess.

This more or less resolves #7344.
We currently don't have support for narrowing based on string or int
literals, but that's a separate issue and should be resolved by
#7169 (which I resumed work
on earlier this week).

Author: Michael0x2a

mypy/checker.py

@@ -5,7 +5,8 @@
 from contextlib import contextmanager
 
 from typing import (
-    Dict, Set, List, cast, Tuple, TypeVar, Union, Optional, NamedTuple, Iterator, Sequence
+    Dict, Set, List, cast, Tuple, TypeVar, Union, Optional, NamedTuple, Iterator, Sequence,
+    Mapping,
 )
 from typing_extensions import Final
 
@@ -42,12 +43,14 @@
 )
 import mypy.checkexpr
 from mypy.checkmember import (
-    analyze_descriptor_access, type_object_type,
+    analyze_member_access, analyze_descriptor_access, type_object_type,
 )
 from mypy.typeops import (
     map_type_from_supertype, bind_self, erase_to_bound, make_simplified_union,
-    erase_def_to_union_or_bound, erase_to_union_or_bound,
-    true_only, false_only, function_type, TypeVarExtractor
+    erase_def_to_union_or_bound, erase_to_union_or_bound, coerce_to_literal,
+    try_getting_str_literals_from_type, try_getting_int_literals_from_type,
+    tuple_fallback, is_singleton_type, try_expanding_enum_to_union,
+    true_only, false_only, function_type, TypeVarExtractor,
 )
 from mypy import message_registry
 from mypy.subtypes import (
@@ -71,9 +74,6 @@
 from mypy.plugin import Plugin, CheckerPluginInterface
 from mypy.sharedparse import BINARY_MAGIC_METHODS
 from mypy.scope import Scope
-from mypy.typeops import (
-    tuple_fallback, coerce_to_literal, is_singleton_type, try_expanding_enum_to_union
-)
 from mypy import state, errorcodes as codes
 from mypy.traverser import has_return_statement, all_return_statements
 from mypy.errorcodes import ErrorCode
@@ -3708,6 +3708,12 @@ def find_isinstance_check(self, node: Expression
 
         Guaranteed to not return None, None. (But may return {}, {})
         """
+        if_map, else_map = self.find_isinstance_check_helper(node)
+        new_if_map = self.propagate_up_typemap_info(self.type_map, if_map)
+        new_else_map = self.propagate_up_typemap_info(self.type_map, else_map)
+        return new_if_map, new_else_map
+
+    def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeMap]:
         type_map = self.type_map
         if is_true_literal(node):
             return {}, None
@@ -3834,28 +3840,196 @@ def find_isinstance_check(self, node: Expression
                         else None)
             return if_map, else_map
         elif isinstance(node, OpExpr) and node.op == 'and':
-            left_if_vars, left_else_vars = self.find_isinstance_check(node.left)
-            right_if_vars, right_else_vars = self.find_isinstance_check(node.right)
+            left_if_vars, left_else_vars = self.find_isinstance_check_helper(node.left)
+            right_if_vars, right_else_vars = self.find_isinstance_check_helper(node.right)
 
             # (e1 and e2) is true if both e1 and e2 are true,
             # and false if at least one of e1 and e2 is false.
             return (and_conditional_maps(left_if_vars, right_if_vars),
                     or_conditional_maps(left_else_vars, right_else_vars))
         elif isinstance(node, OpExpr) and node.op == 'or':
-            left_if_vars, left_else_vars = self.find_isinstance_check(node.left)
-            right_if_vars, right_else_vars = self.find_isinstance_check(node.right)
+            left_if_vars, left_else_vars = self.find_isinstance_check_helper(node.left)
+            right_if_vars, right_else_vars = self.find_isinstance_check_helper(node.right)
 
             # (e1 or e2) is true if at least one of e1 or e2 is true,
             # and false if both e1 and e2 are false.
             return (or_conditional_maps(left_if_vars, right_if_vars),
                     and_conditional_maps(left_else_vars, right_else_vars))
         elif isinstance(node, UnaryExpr) and node.op == 'not':
-            left, right = self.find_isinstance_check(node.expr)
+            left, right = self.find_isinstance_check_helper(node.expr)
             return right, left
 
         # Not a supported isinstance check
         return {}, {}
 
+    def propagate_up_typemap_info(self,
+                                  existing_types: Mapping[Expression, Type],
+                                  new_types: TypeMap) -> TypeMap:
+        """Attempts refining parent expressions of any MemberExpr or IndexExprs in new_types.
+
+        Specifically, this function accepts two mappings of expression to original types:
+        the original mapping (existing_types), and a new mapping (new_types) intended to
+        update the original.
+
+        This function iterates through new_types and attempts to use the information to try
+        refining any parent types that happen to be unions.
+
+        For example, suppose there are two types "A = Tuple[int, int]" and "B = Tuple[str, str]".
+        Next, suppose that 'new_types' specifies the expression 'foo[0]' has a refined type
+        of 'int' and that 'foo' was previously deduced to be of type Union[A, B].
+
+        Then, this function will observe that since A[0] is an int and B[0] is not, the type of
+        'foo' can be further refined from Union[A, B] into just B.
+
+        We perform this kind of "parent narrowing" for member lookup expressions and indexing
+        expressions into tuples, namedtuples, and typeddicts. We repeat this narrowing
+        recursively if the parent is also a "lookup expression". So for example, if we have
+        the expression "foo['bar'].baz[0]", we'd potentially end up refining types for the
+        expressions "foo", "foo['bar']", and "foo['bar'].baz".
+
+        We return the newly refined map. This map is guaranteed to be a superset of 'new_types'.
+        """
+        if new_types is None:
+            return None
+        output_map = {}
+        for expr, expr_type in new_types.items():
+            # The original inferred type should always be present in the output map, of course
+            output_map[expr] = expr_type
+
+            # Next, try using this information to refine the parent types, if applicable.
+            new_mapping = self.refine_parent_types(existing_types, expr, expr_type)
+            for parent_expr, proposed_parent_type in new_mapping.items():
+                # We don't try inferring anything if we've already inferred something for
+                # the parent expression.
+                # TODO: Consider picking the narrower type instead of always discarding this?
+                if parent_expr in new_types:
+                    continue
+                output_map[parent_expr] = proposed_parent_type
+        return output_map
+
+    def refine_parent_types(self,
+                            existing_types: Mapping[Expression, Type],
+                            expr: Expression,
+                            expr_type: Type) -> Mapping[Expression, Type]:
+        """Checks if the given expr is a 'lookup operation' into a union and iteratively refines
+        the parent types based on the 'expr_type'.
+
+        For example, if 'expr' is an expression like 'a.b.c.d', we'll potentially return refined
+        types for expressions 'a', 'a.b', and 'a.b.c'.
+
+        For more details about what a 'lookup operation' is and how we use the expr_type to refine
+        the parent types of lookup_expr, see the docstring in 'propagate_up_typemap_info'.
+        """
+        output = {}  # type: Dict[Expression, Type]
+
+        # Note: parent_expr and parent_type are progressively refined as we crawl up the
+        # parent lookup chain.
+        while True:
+            # First, check if this expression is one that's attempting to
+            # "lookup" some key in the parent type. If so, save the parent type
+            # and create function that will try replaying the same lookup
+            # operation against arbitrary types.
+            if isinstance(expr, MemberExpr):
+                parent_expr = expr.expr
+                parent_type = existing_types.get(parent_expr)
+                member_name = expr.name
+
+                def replay_lookup(new_parent_type: ProperType) -> Optional[Type]:
+                    msg_copy = self.msg.clean_copy()
+                    msg_copy.disable_count = 0
+                    member_type = analyze_member_access(
+                        name=member_name,
+                        typ=new_parent_type,
+                        context=parent_expr,
+                        is_lvalue=False,
+                        is_super=False,
+                        is_operator=False,
+                        msg=msg_copy,
+                        original_type=new_parent_type,
+                        chk=self,
+                        in_literal_context=False,
+                    )
+                    if msg_copy.is_errors():
+                        return None
+                    else:
+                        return member_type
+            elif isinstance(expr, IndexExpr):
+                parent_expr = expr.base
+                parent_type = existing_types.get(parent_expr)
+
+                index_type = existing_types.get(expr.index)
+                if index_type is None:
+                    return output
+
+                str_literals = try_getting_str_literals_from_type(index_type)
+                if str_literals is not None:
+                    # Refactoring these two indexing replay functions is surprisingly
+                    # tricky -- see https://github.com/python/mypy/pull/7917, which
+                    # was blocked by https://github.com/mypyc/mypyc/issues/586
+                    def replay_lookup(new_parent_type: ProperType) -> Optional[Type]:
+                        if not isinstance(new_parent_type, TypedDictType):
+                            return None
+                        try:
+                            assert str_literals is not None
+                            member_types = [new_parent_type.items[key] for key in str_literals]
+                        except KeyError:
+                            return None
+                        return make_simplified_union(member_types)
+                else:
+                    int_literals = try_getting_int_literals_from_type(index_type)
+                    if int_literals is not None:
+                        def replay_lookup(new_parent_type: ProperType) -> Optional[Type]:
+                            if not isinstance(new_parent_type, TupleType):
+                                return None
+                            try:
+                                assert int_literals is not None
+                                member_types = [new_parent_type.items[key] for key in int_literals]
+                            except IndexError:
+                                return None
+                            return make_simplified_union(member_types)
+                    else:
+                        return output
+            else:
+                return output
+
+            # If we somehow didn't previously derive the parent type, abort completely
+            # with what we have so far: something went wrong at an earlier stage.
+            if parent_type is None:
+                return output
+
+            # We currently only try refining the parent type if it's a Union.
+            # If not, there's no point in trying to refine any further parents
+            # since we have no further information we can use to refine the lookup
+            # chain, so we end early as an optimization.
+            parent_type = get_proper_type(parent_type)
+            if not isinstance(parent_type, UnionType):
+                return output
+
+            # Take each element in the parent union and replay the original lookup procedure
+            # to figure out which parents are compatible.
+            new_parent_types = []
+            for item in parent_type.items:
+                item = get_proper_type(item)
+                member_type = replay_lookup(item)
+                if member_type is None:
+                    # We were unable to obtain the member type. So, we give up on refining this
+                    # parent type entirely and abort.
+                    return output
+
+                if is_overlapping_types(member_type, expr_type):
+                    new_parent_types.append(item)
+
+            # If none of the parent types overlap (if we derived an empty union), something
+            # went wrong. We should never hit this case, but deriving the uninhabited type or
+            # reporting an error both seem unhelpful. So we abort.
+            if not new_parent_types:
+                return output
+
+            expr = parent_expr
+            expr_type = output[parent_expr] = make_simplified_union(new_parent_types)
+
+        return output
+
     #
     # Helpers
     #

mypy/checkexpr.py

@@ -2709,6 +2709,9 @@ def visit_index_with_type(self, left_type: Type, e: IndexExpr,
         index = e.index
         left_type = get_proper_type(left_type)
 
+        # Visit the index, just to make sure we have a type for it available
+        self.accept(index)
+
         if isinstance(left_type, UnionType):
             original_type = original_type or left_type
             return make_simplified_union([self.visit_index_with_type(typ, e,

mypy/test/testcheck.py

@@ -46,6 +46,7 @@
     'check-isinstance.test',
     'check-lists.test',
     'check-namedtuple.test',
+    'check-narrowing.test',
     'check-typeddict.test',
     'check-type-aliases.test',
     'check-ignore.test',

mypy/typeops.py

@@ -5,17 +5,18 @@
       since these may assume that MROs are ready.
 """
 
-from typing import cast, Optional, List, Sequence, Set, Iterable
+from typing import cast, Optional, List, Sequence, Set, Iterable, TypeVar
+from typing_extensions import Type as TypingType
 import sys
 
 from mypy.types import (
     TupleType, Instance, FunctionLike, Type, CallableType, TypeVarDef, Overloaded,
-    TypeVarType, UninhabitedType, FormalArgument, UnionType, NoneType,
+    TypeVarType, UninhabitedType, FormalArgument, UnionType, NoneType, TypedDictType,
     AnyType, TypeOfAny, TypeType, ProperType, LiteralType, get_proper_type, get_proper_types,
     copy_type, TypeAliasType, TypeQuery
 )
 from mypy.nodes import (
-    FuncBase, FuncItem, OverloadedFuncDef, TypeInfo, TypeVar, ARG_STAR, ARG_STAR2, ARG_POS,
+    FuncBase, FuncItem, OverloadedFuncDef, TypeInfo, ARG_STAR, ARG_STAR2, ARG_POS,
     Expression, StrExpr, Var
 )
 from mypy.maptype import map_instance_to_supertype
@@ -43,6 +44,25 @@ def tuple_fallback(typ: TupleType) -> Instance:
     return Instance(info, [join_type_list(typ.items)])
 
 
+def try_getting_instance_fallback(typ: ProperType) -> Optional[Instance]:
+    """Returns the Instance fallback for this type if one exists.
+
+    Otherwise, returns None.
+    """
+    if isinstance(typ, Instance):
+        return typ
+    elif isinstance(typ, TupleType):
+        return tuple_fallback(typ)
+    elif isinstance(typ, TypedDictType):
+        return typ.fallback
+    elif isinstance(typ, FunctionLike):
+        return typ.fallback
+    elif isinstance(typ, LiteralType):
+        return typ.fallback
+    else:
+        return None
+
+
 def type_object_type_from_function(signature: FunctionLike,
                                    info: TypeInfo,
                                    def_info: TypeInfo,
@@ -481,27 +501,66 @@ def try_getting_str_literals(expr: Expression, typ: Type) -> Optional[List[str]]
     2. 'typ' is a LiteralType containing a string
     3. 'typ' is a UnionType containing only LiteralType of strings
     """
-    typ = get_proper_type(typ)
-
     if isinstance(expr, StrExpr):
         return [expr.value]
 
+    # TODO: See if we can eliminate this function and call the below one directly
+    return try_getting_str_literals_from_type(typ)
+
+
+def try_getting_str_literals_from_type(typ: Type) -> Optional[List[str]]:
+    """If the given expression or type corresponds to a string Literal
+    or a union of string Literals, returns a list of the underlying strings.
+    Otherwise, returns None.
+
+    For example, if we had the type 'Literal["foo", "bar"]' as input, this function
+    would return a list of strings ["foo", "bar"].
+    """
+    return try_getting_literals_from_type(typ, str, "builtins.str")
+
+
+def try_getting_int_literals_from_type(typ: Type) -> Optional[List[int]]:
+    """If the given expression or type corresponds to an int Literal
+    or a union of int Literals, returns a list of the underlying ints.
+    Otherwise, returns None.
+
+    For example, if we had the type 'Literal[1, 2, 3]' as input, this function
+    would return a list of ints [1, 2, 3].
+    """
+    return try_getting_literals_from_type(typ, int, "builtins.int")
+
+
+T = TypeVar('T')
+
+
+def try_getting_literals_from_type(typ: Type,
+                                   target_literal_type: TypingType[T],
+                                   target_fullname: str) -> Optional[List[T]]:
+    """If the given expression or type corresponds to a Literal or
+    union of Literals where the underlying values corresponds to the given
+    target type, returns a list of those underlying values. Otherwise,
+    returns None.
+    """
+    typ = get_proper_type(typ)
+
     if isinstance(typ, Instance) and typ.last_known_value is not None:
         possible_literals = [typ.last_known_value]  # type: List[Type]
     elif isinstance(typ, UnionType):
         possible_literals = list(typ.items)
     else:
         possible_literals = [typ]
 
-    strings = []
+    literals = []  # type: List[T]
     for lit in get_proper_types(possible_literals):
-        if isinstance(lit, LiteralType) and lit.fallback.type.fullname() == 'builtins.str':
+        if isinstance(lit, LiteralType) and lit.fallback.type.fullname() == target_fullname:
             val = lit.value
-            assert isinstance(val, str)
-            strings.append(val)
+            if isinstance(val, target_literal_type):
+                literals.append(val)
+            else:
+                return None
         else:
             return None
-    return strings
+    return literals
 
 
 def get_enum_values(typ: Instance) -> List[str]: