Add support for narrowing Literals using equality

This pull request (finally) adds support for narrowing expressions
using Literal types by equality, instead of just identity. For example,
the following "tagged union" pattern is now supported:

```python
class Foo(TypedDict):
    key: Literal["A"]
    blah: int

class Bar(TypedDict):
    key: Literal["B"]
    something: str

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

Previously, this was possible to do only with Enum Literals and the
`is` operator, which is perhaps not very intuitive.

The main limitation with this pull request is that it'll perform narrowing
only if either the LHS or RHS contains an explicit Literal type
somewhere. If this limitation is not present, we end up breaking a decent
amount of real-world code -- mostly tests -- that do something like this:

```python
def some_test_case() -> None:
    worker = Worker()

    # Without the limitation, we narrow 'worker.state' to
    # Literal['ready'] in this assert...
    assert worker.state == 'ready'

    worker.start()

    # ...which subsequently causes this second assert to narrow
    # worker.state to <uninhabited>, causing the last line to be
    # unreachable.
    assert worker.state == 'running'
    worker.query()
```

I tried for several weeks to find a more intelligent way around this
problem, but everything I tried ended up being either insufficient or
super-hacky, so I gave up and went for this brute-force solution.

The other main limitation is that we perform narrowing only if both the
LHS and RHS do not define custom `__eq__` or `__ne__` methods, but this
seems like a more reasonable one to me.

Resolves #7944.

Author: Michael0x2a

mypy/checker.py

@@ -6,7 +6,7 @@
 
 from typing import (
     Dict, Set, List, cast, Tuple, TypeVar, Union, Optional, NamedTuple, Iterator, Iterable,
-    Sequence, Mapping, Generic, AbstractSet
+    Sequence, Mapping, Generic, AbstractSet, Callable
 )
 from typing_extensions import Final
 
@@ -50,7 +50,8 @@
     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,
+    true_only, false_only, function_type, TypeVarExtractor, custom_special_method,
+    is_literal_type_like,
 )
 from mypy import message_registry
 from mypy.subtypes import (
@@ -3890,20 +3891,59 @@ def find_isinstance_check_helper(self, node: Expression) -> Tuple[TypeMap, TypeM
 
             partial_type_maps = []
             for operator, expr_indices in simplified_operator_list:
-                if operator in {'is', 'is not'}:
-                    if_map, else_map = self.refine_identity_comparison_expression(
-                        operands,
-                        operand_types,
-                        expr_indices,
-                        narrowable_operand_index_to_hash.keys(),
-                    )
-                elif operator in {'==', '!='}:
-                    if_map, else_map = self.refine_equality_comparison_expression(
-                        operands,
-                        operand_types,
-                        expr_indices,
-                        narrowable_operand_index_to_hash.keys(),
-                    )
+                if operator in {'is', 'is not', '==', '!='}:
+                    # is_valid_target:
+                    #   Controls which types we're allowed to narrow exprs to. Note that
+                    #   we cannot use 'is_literal_type_like' in both cases since doing
+                    #   'x = 10000 + 1; x is 10001' is not always True in all Python impls.
+                    #
+                    # coerce_only_in_literal_context:
+                    #   If true, coerce types into literal types only if one or more of
+                    #   the provided exprs contains an explicit Literal type. This could
+                    #   technically be set to any arbitrary value, but it seems being liberal
+                    #   with narrowing when using 'is' and conservative when using '==' seems
+                    #   to break the least amount of real-world code.
+                    #
+                    # should_narrow_by_identity:
+                    #   Set to 'false' only if the user defines custom __eq__ or __ne__ methods
+                    #   that could cause identity-based narrowing to produce invalid results.
+                    if operator in {'is', 'is not'}:
+                        is_valid_target = is_singleton_type    # type: Callable[[Type], bool]
+                        coerce_only_in_literal_context = False
+                        should_narrow_by_identity = True
+                    else:
+                        is_valid_target = is_exactly_literal_type
+                        coerce_only_in_literal_context = True
+
+                        def has_no_custom_eq_checks(t: Type) -> bool:
+                            return not custom_special_method(t, '__eq__', check_all=False) \
+                                   and not custom_special_method(t, '__ne__', check_all=False)
+                        expr_types = [operand_types[i] for i in expr_indices]
+                        should_narrow_by_identity = all(map(has_no_custom_eq_checks, expr_types))
+
+                    if_map = {}    # type: TypeMap
+                    else_map = {}  # type: TypeMap
+                    if should_narrow_by_identity:
+                        if_map, else_map = self.refine_identity_comparison_expression(
+                            operands,
+                            operand_types,
+                            expr_indices,
+                            narrowable_operand_index_to_hash.keys(),
+                            is_valid_target,
+                            coerce_only_in_literal_context,
+                        )
+
+                    # Strictly speaking, we should also skip this check if the objects in the expr
+                    # chain have custom __eq__ or __ne__ methods. But we (maybe optimistically)
+                    # assume nobody would actually create a custom objects that considers itself
+                    # equal to None.
+                    if if_map == {} and else_map == {}:
+                        if_map, else_map = self.refine_away_none_in_comparison(
+                            operands,
+                            operand_types,
+                            expr_indices,
+                            narrowable_operand_index_to_hash.keys(),
+                        )
                 elif operator in {'in', 'not in'}:
                     assert len(expr_indices) == 2
                     left_index, right_index = expr_indices
@@ -4146,8 +4186,10 @@ def refine_identity_comparison_expression(self,
                                               operand_types: List[Type],
                                               chain_indices: List[int],
                                               narrowable_operand_indices: AbstractSet[int],
+                                              is_valid_target: Callable[[ProperType], bool],
+                                              coerce_only_in_literal_context: bool,
                                               ) -> Tuple[TypeMap, TypeMap]:
-        """Produces conditional type maps refining expressions used in an identity comparison.
+        """Produces conditional type maps refining exprs used in an identity/equality comparison.
 
         The 'operands' and 'operand_types' lists should be the full list of operands used
         in the overall comparison expression. The 'chain_indices' list is the list of indices
@@ -4163,30 +4205,45 @@ def refine_identity_comparison_expression(self,
         The 'narrowable_operand_indices' parameter is the set of all indices we are allowed
         to refine the types of: that is, all operands that will potentially be a part of
         the output TypeMaps.
+
+        Although this function could theoretically try setting the types of the operands
+        in the chains to the meet, doing that causes too many issues in real-world code.
+        Instead, we use 'is_valid_target' to identify which of the given chain types
+        we could plausibly use as the refined type for the expressions in the chain.
+
+        Similarly, 'coerce_only_in_literal_context' controls whether we should try coercing
+        expressions in the chain to a Literal type. Performing this coercion is sometimes
+        too aggressive of a narrowing, depending on context.
         """
-        singleton = None  # type: Optional[ProperType]
-        possible_singleton_indices = []
+        should_coerce = True
+        if coerce_only_in_literal_context:
+            should_coerce = any(is_literal_type_like(operand_types[i]) for i in chain_indices)
+
+        target = None  # type: Optional[Type]
+        possible_target_indices = []
         for i in chain_indices:
-            coerced_type = coerce_to_literal(operand_types[i])
-            if not is_singleton_type(coerced_type):
+            expr_type = operand_types[i]
+            if should_coerce:
+                expr_type = coerce_to_literal(expr_type)
+            if not is_valid_target(get_proper_type(expr_type)):
                 continue
-            if singleton and not is_same_type(singleton, coerced_type):
-                # We have multiple disjoint singleton types. So the 'if' branch
+            if target and not is_same_type(target, expr_type):
+                # We have multiple disjoint target types. So the 'if' branch
                 # must be unreachable.
                 return None, {}
-            singleton = coerced_type
-            possible_singleton_indices.append(i)
+            target = expr_type
+            possible_target_indices.append(i)
 
-        # There's nothing we can currently infer if none of the operands are singleton types,
+        # There's nothing we can currently infer if none of the operands are valid targets,
         # so we end early and infer nothing.
-        if singleton is None:
+        if target is None:
             return {}, {}
 
-        # If possible, use an unassignable expression as the singleton.
-        # We skip refining the type of the singleton below, so ideally we'd
+        # If possible, use an unassignable expression as the target.
+        # We skip refining the type of the target below, so ideally we'd
         # want to pick an expression we were going to skip anyways.
         singleton_index = -1
-        for i in possible_singleton_indices:
+        for i in possible_target_indices:
             if i not in narrowable_operand_indices:
                 singleton_index = i
 
@@ -4215,20 +4272,21 @@ def refine_identity_comparison_expression(self,
         # currently will just mark the whole branch as unreachable if either operand is
         # narrowed to <uninhabited>.
         if singleton_index == -1:
-            singleton_index = possible_singleton_indices[-1]
+            singleton_index = possible_target_indices[-1]
 
         enum_name = None
-        if isinstance(singleton, LiteralType) and singleton.is_enum_literal():
-            enum_name = singleton.fallback.type.fullname
+        target = get_proper_type(target)
+        if isinstance(target, LiteralType) and target.is_enum_literal():
+            enum_name = target.fallback.type.fullname
 
-        target_type = [TypeRange(singleton, is_upper_bound=False)]
+        target_type = [TypeRange(target, is_upper_bound=False)]
 
         partial_type_maps = []
         for i in chain_indices:
-            # If we try refining a singleton against itself, conditional_type_map
+            # If we try refining a type against itself, conditional_type_map
             # will end up assuming that the 'else' branch is unreachable. This is
             # typically not what we want: generally the user will intend for the
-            # singleton type to be some fixed 'sentinel' value and will want to refine
+            # target type to be some fixed 'sentinel' value and will want to refine
             # the other exprs against this one instead.
             if i == singleton_index:
                 continue
@@ -4246,17 +4304,16 @@ def refine_identity_comparison_expression(self,
 
         return reduce_partial_conditional_maps(partial_type_maps)
 
-    def refine_equality_comparison_expression(self,
-                                              operands: List[Expression],
-                                              operand_types: List[Type],
-                                              chain_indices: List[int],
-                                              narrowable_operand_indices: AbstractSet[int],
-                                              ) -> Tuple[TypeMap, TypeMap]:
-        """Produces conditional type maps refining expressions used in an equality comparison.
+    def refine_away_none_in_comparison(self,
+                                       operands: List[Expression],
+                                       operand_types: List[Type],
+                                       chain_indices: List[int],
+                                       narrowable_operand_indices: AbstractSet[int],
+                                       ) -> Tuple[TypeMap, TypeMap]:
+        """Produces conditional type maps refining away None in an identity/equality chain.
 
-        For more details, see the docstring of 'refine_equality_comparison' up above.
-        The only difference is that this function is for refining equality operations
-        (e.g. 'a == b == c') instead of identity ('a is b is c').
+        For more details about what the different arguments mean, see the
+        docstring of 'refine_identity_comparison_expression' up above.
         """
         non_optional_types = []
         for i in chain_indices:
@@ -4749,7 +4806,7 @@ class Foo(Enum):
         return False
 
     parent_type = get_proper_type(parent_type)
-    member_type = coerce_to_literal(member_type)
+    member_type = get_proper_type(coerce_to_literal(member_type))
     if not isinstance(parent_type, FunctionLike) or not isinstance(member_type, LiteralType):
         return False
 
@@ -5540,3 +5597,9 @@ def has_bool_item(typ: ProperType) -> bool:
         return any(is_named_instance(item, 'builtins.bool')
                    for item in typ.items)
     return False
+
+
+# TODO: why can't we define this as an inline function?
+# Does mypyc not support them?
+def is_exactly_literal_type(t: Type) -> bool:
+    return isinstance(get_proper_type(t), LiteralType)

mypy/checkexpr.py

@@ -32,7 +32,6 @@
     YieldFromExpr, TypedDictExpr, PromoteExpr, NewTypeExpr, NamedTupleExpr, TypeVarExpr,
     TypeAliasExpr, BackquoteExpr, EnumCallExpr, TypeAlias, SymbolNode, PlaceholderNode,
     ARG_POS, ARG_OPT, ARG_NAMED, ARG_STAR, ARG_STAR2, LITERAL_TYPE, REVEAL_TYPE,
-    SYMBOL_FUNCBASE_TYPES
 )
 from mypy.literals import literal
 from mypy import nodes
@@ -51,15 +50,16 @@
 from mypy import erasetype
 from mypy.checkmember import analyze_member_access, type_object_type
 from mypy.argmap import ArgTypeExpander, map_actuals_to_formals, map_formals_to_actuals
-from mypy.checkstrformat import StringFormatterChecker, custom_special_method
+from mypy.checkstrformat import StringFormatterChecker
 from mypy.expandtype import expand_type, expand_type_by_instance, freshen_function_type_vars
 from mypy.util import split_module_names
 from mypy.typevars import fill_typevars
 from mypy.visitor import ExpressionVisitor
 from mypy.plugin import Plugin, MethodContext, MethodSigContext, FunctionContext
 from mypy.typeops import (
     tuple_fallback, make_simplified_union, true_only, false_only, erase_to_union_or_bound,
-    function_type, callable_type, try_getting_str_literals
+    function_type, callable_type, try_getting_str_literals, custom_special_method,
+    is_literal_type_like,
 )
 import mypy.errorcodes as codes
 
@@ -4266,24 +4266,6 @@ def merge_typevars_in_callables_by_name(
     return output, variables
 
 
-def is_literal_type_like(t: Optional[Type]) -> bool:
-    """Returns 'true' if the given type context is potentially either a LiteralType,
-    a Union of LiteralType, or something similar.
-    """
-    t = get_proper_type(t)
-    if t is None:
-        return False
-    elif isinstance(t, LiteralType):
-        return True
-    elif isinstance(t, UnionType):
-        return any(is_literal_type_like(item) for item in t.items)
-    elif isinstance(t, TypeVarType):
-        return (is_literal_type_like(t.upper_bound)
-                or any(is_literal_type_like(item) for item in t.values))
-    else:
-        return False
-
-
 def try_getting_literal(typ: Type) -> ProperType:
     """If possible, get a more precise literal type for a given type."""
     typ = get_proper_type(typ)
@@ -4305,29 +4287,6 @@ def is_expr_literal_type(node: Expression) -> bool:
     return False
 
 
-def custom_equality_method(typ: Type) -> bool:
-    """Does this type have a custom __eq__() method?"""
-    typ = get_proper_type(typ)
-    if isinstance(typ, Instance):
-        method = typ.type.get('__eq__')
-        if method and isinstance(method.node, (SYMBOL_FUNCBASE_TYPES, Decorator, Var)):
-            if method.node.info:
-                return not method.node.info.fullname.startswith('builtins.')
-        return False
-    if isinstance(typ, UnionType):
-        return any(custom_equality_method(t) for t in typ.items)
-    if isinstance(typ, TupleType):
-        return custom_equality_method(tuple_fallback(typ))
-    if isinstance(typ, CallableType) and typ.is_type_obj():
-        # Look up __eq__ on the metaclass for class objects.
-        return custom_equality_method(typ.fallback)
-    if isinstance(typ, AnyType):
-        # Avoid false positives in uncertain cases.
-        return True
-    # TODO: support other types (see ExpressionChecker.has_member())?
-    return False
-
-
 def has_bytes_component(typ: Type, py2: bool = False) -> bool:
     """Is this one of builtin byte types, or a union that contains it?"""
     typ = get_proper_type(typ)

mypy/checkstrformat.py

@@ -19,12 +19,12 @@
 
 from mypy.types import (
     Type, AnyType, TupleType, Instance, UnionType, TypeOfAny, get_proper_type, TypeVarType,
-    CallableType, LiteralType, get_proper_types
+    LiteralType, get_proper_types
 )
 from mypy.nodes import (
     StrExpr, BytesExpr, UnicodeExpr, TupleExpr, DictExpr, Context, Expression, StarExpr, CallExpr,
     IndexExpr, MemberExpr, TempNode, ARG_POS, ARG_STAR, ARG_NAMED, ARG_STAR2,
-    SYMBOL_FUNCBASE_TYPES, Decorator, Var, Node, MypyFile, ExpressionStmt, NameExpr, IntExpr
+    Node, MypyFile, ExpressionStmt, NameExpr, IntExpr
 )
 import mypy.errorcodes as codes
 
@@ -35,7 +35,7 @@
 from mypy import message_registry
 from mypy.messages import MessageBuilder
 from mypy.maptype import map_instance_to_supertype
-from mypy.typeops import tuple_fallback
+from mypy.typeops import custom_special_method
 from mypy.subtypes import is_subtype
 from mypy.parse import parse
 
@@ -961,32 +961,3 @@ def has_type_component(typ: Type, fullname: str) -> bool:
     elif isinstance(typ, UnionType):
         return any(has_type_component(t, fullname) for t in typ.relevant_items())
     return False
-
-
-def custom_special_method(typ: Type, name: str,
-                          check_all: bool = False) -> bool:
-    """Does this type have a custom special method such as __format__() or __eq__()?
-
-    If check_all is True ensure all items of a union have a custom method, not just some.
-    """
-    typ = get_proper_type(typ)
-    if isinstance(typ, Instance):
-        method = typ.type.get(name)
-        if method and isinstance(method.node, (SYMBOL_FUNCBASE_TYPES, Decorator, Var)):
-            if method.node.info:
-                return not method.node.info.fullname.startswith('builtins.')
-        return False
-    if isinstance(typ, UnionType):
-        if check_all:
-            return all(custom_special_method(t, name, check_all) for t in typ.items)
-        return any(custom_special_method(t, name) for t in typ.items)
-    if isinstance(typ, TupleType):
-        return custom_special_method(tuple_fallback(typ), name)
-    if isinstance(typ, CallableType) and typ.is_type_obj():
-        # Look up __method__ on the metaclass for class objects.
-        return custom_special_method(typ.fallback, name)
-    if isinstance(typ, AnyType):
-        # Avoid false positives in uncertain cases.
-        return True
-    # TODO: support other types (see ExpressionChecker.has_member())?
-    return False