Better handling of generics when narrowing (#20863)

Notably we preserve behaviour on the `testNarrowingCollections` test I
added in a previous PR

Closes #20673

Author: hauntsaninja

mypy/checker.py

@@ -32,7 +32,12 @@
 )
 from mypy.checkpattern import PatternChecker
 from mypy.constraints import SUPERTYPE_OF
-from mypy.erasetype import erase_type, erase_typevars, remove_instance_last_known_values
+from mypy.erasetype import (
+    erase_type,
+    erase_typevars,
+    remove_instance_last_known_values,
+    shallow_erase_type_for_equality,
+)
 from mypy.errorcodes import TYPE_VAR, UNUSED_AWAITABLE, UNUSED_COROUTINE, ErrorCode
 from mypy.errors import (
     ErrorInfo,
@@ -6628,6 +6633,9 @@ def comparison_type_narrowing_helper(self, node: ComparisonExpr) -> tuple[TypeMa
                             narrowable_indices={0},
                         )
 
+                        # TODO: This remove_optional code should no longer be needed. The only
+                        # thing it does is paper over a pre-existing deficiency in equality
+                        # narrowing w.r.t to enums.
                         # We only try and narrow away 'None' for now
                         if (
                             not is_unreachable_map(if_map)
@@ -6775,8 +6783,7 @@ def narrow_type_by_identity_equality(
                 target = TypeRange(target_type, is_upper_bound=False)
 
                 if_map, else_map = conditional_types_to_typemaps(
-                    operands[i],
-                    *conditional_types(expr_type, [target], consider_promotion_overlap=True),
+                    operands[i], *conditional_types(expr_type, [target], from_equality=True)
                 )
                 if is_target_for_value_narrowing(get_proper_type(target_type)):
                     all_if_maps.append(if_map)
@@ -6814,9 +6821,7 @@ def narrow_type_by_identity_equality(
                     if is_target_for_value_narrowing(get_proper_type(target_type)):
                         if_map, else_map = conditional_types_to_typemaps(
                             operands[i],
-                            *conditional_types(
-                                expr_type, [target], consider_promotion_overlap=True
-                            ),
+                            *conditional_types(expr_type, [target], from_equality=True),
                         )
                         all_else_maps.append(else_map)
                 continue
@@ -6855,7 +6860,7 @@ def narrow_type_by_identity_equality(
                     if_map, else_map = conditional_types_to_typemaps(
                         operands[i],
                         *conditional_types(
-                            expr_type, [target], default=expr_type, consider_promotion_overlap=True
+                            expr_type, [target], default=expr_type, from_equality=True
                         ),
                     )
                     or_if_maps.append(if_map)
@@ -8359,7 +8364,7 @@ def conditional_types(
     default: None = None,
     *,
     consider_runtime_isinstance: bool = True,
-    consider_promotion_overlap: bool = False,
+    from_equality: bool = False,
 ) -> tuple[Type | None, Type | None]: ...
 
 
@@ -8370,7 +8375,7 @@ def conditional_types(
     default: Type,
     *,
     consider_runtime_isinstance: bool = True,
-    consider_promotion_overlap: bool = False,
+    from_equality: bool = False,
 ) -> tuple[Type, Type]: ...
 
 
@@ -8380,7 +8385,7 @@ def conditional_types(
     default: Type | None = None,
     *,
     consider_runtime_isinstance: bool = True,
-    consider_promotion_overlap: bool = False,
+    from_equality: bool = False,
 ) -> tuple[Type | None, Type | None]:
     """Takes in the current type and a proposed type of an expression.
 
@@ -8425,7 +8430,7 @@ def conditional_types(
                 proposed_type_ranges,
                 default=union_item,
                 consider_runtime_isinstance=consider_runtime_isinstance,
-                consider_promotion_overlap=consider_promotion_overlap,
+                from_equality=from_equality,
             )
             yes_items.append(yes_type)
             no_items.append(no_type)
@@ -8470,17 +8475,23 @@ def conditional_types(
                 consider_runtime_isinstance=consider_runtime_isinstance,
             )
             return default, remainder
-    if not is_overlapping_types(
-        current_type, proposed_type, ignore_promotions=not consider_promotion_overlap
-    ):
-        # Expression is never of any type in proposed_type_ranges
-        return UninhabitedType(), default
-    if consider_promotion_overlap and not is_overlapping_types(
-        current_type, proposed_type, ignore_promotions=True
-    ):
-        # We set consider_promotion_overlap when comparing equality. This is one of the places
-        # at runtime where subtyping with promotion does happen to match runtime semantics
-        return default, default
+
+    if from_equality:
+        # We erase generic args because values with different generic types can compare equal
+        # For instance, cast(list[str], []) and cast(list[int], [])
+        proposed_type = shallow_erase_type_for_equality(proposed_type)
+        if not is_overlapping_types(current_type, proposed_type, ignore_promotions=False):
+            # Equality narrowing is one of the places at runtime where subtyping with promotion
+            # does happen to match runtime semantics
+            # Expression is never of any type in proposed_type_ranges
+            return UninhabitedType(), default
+        if not is_overlapping_types(current_type, proposed_type, ignore_promotions=True):
+            return default, default
+    else:
+        if not is_overlapping_types(current_type, proposed_type, ignore_promotions=True):
+            # Expression is never of any type in proposed_type_ranges
+            return UninhabitedType(), default
+
     # we can only restrict when the type is precise, not bounded
     proposed_precise_type = UnionType.make_union(
         [type_range.item for type_range in proposed_type_ranges if not type_range.is_upper_bound]
@@ -8726,13 +8737,7 @@ def reduce_and_conditional_type_maps(ms: list[TypeMap], *, use_meet: bool) -> Ty
     return result
 
 
-BUILTINS_CUSTOM_EQ_CHECKS: Final = {
-    "builtins.bytearray",
-    "builtins.memoryview",
-    "builtins.list",
-    "builtins.dict",
-    "builtins.set",
-}
+BUILTINS_CUSTOM_EQ_CHECKS: Final = {"builtins.bytearray", "builtins.memoryview"}
 
 
 def has_custom_eq_checks(t: Type) -> bool:

mypy/erasetype.py

@@ -285,3 +285,17 @@ def visit_union_type(self, t: UnionType) -> Type:
                     merged.append(orig_item)
             return UnionType.make_union(merged)
         return new
+
+
+def shallow_erase_type_for_equality(typ: Type) -> ProperType:
+    """Erase type variables from Instance's"""
+    p_typ = get_proper_type(typ)
+    if isinstance(p_typ, Instance):
+        if not p_typ.args:
+            return p_typ
+        args = erased_vars(p_typ.type.defn.type_vars, TypeOfAny.special_form)
+        return Instance(p_typ.type, args, p_typ.line)
+    if isinstance(p_typ, UnionType):
+        items = [shallow_erase_type_for_equality(item) for item in p_typ.items]
+        return UnionType.make_union(items)
+    return p_typ

test-data/unit/check-narrowing.test

@@ -1065,6 +1065,92 @@ def f(x: Custom, y: CustomSub):
         reveal_type(y)  # N: Revealed type is "__main__.CustomSub"
 [builtins fixtures/tuple.pyi]
 
+[case testNarrowingCustomEqualityGeneric]
+# flags: --strict-equality --warn-unreachable
+from __future__ import annotations
+from typing import Union
+
+class Custom:
+    def __eq__(self, other: object) -> bool:
+        raise
+
+class Default: ...
+
+def f1(x: list[Custom] | Default, y: list[int]):
+    if x == y:  # E: Non-overlapping equality check (left operand type: "list[Custom] | Default", right operand type: "list[int]")
+        reveal_type(x)  # N: Revealed type is "builtins.list[__main__.Custom]"
+        reveal_type(y)  # N: Revealed type is "builtins.list[builtins.int]"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.list[__main__.Custom] | __main__.Default"
+        reveal_type(y)  # N: Revealed type is "builtins.list[builtins.int]"
+
+f1([], [])
+
+def f2(x: list[Custom] | Default, y: list[int] | list[Default]):
+    if x == y:  # E: Non-overlapping equality check (left operand type: "list[Custom] | Default", right operand type: "list[int] | list[Default]")
+        reveal_type(x)  # N: Revealed type is "builtins.list[__main__.Custom]"
+        reveal_type(y)  # N: Revealed type is "builtins.list[builtins.int] | builtins.list[__main__.Default]"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.list[__main__.Custom] | __main__.Default"
+        reveal_type(y)  # N: Revealed type is "builtins.list[builtins.int] | builtins.list[__main__.Default]"
+
+listcustom_or_default = Union[list[Custom], Default]
+listint_or_default = Union[list[int], list[Default]]
+
+def f2_with_alias(x: listcustom_or_default, y: listint_or_default):
+    if x == y:  # E: Non-overlapping equality check (left operand type: "list[Custom] | Default", right operand type: "list[int] | list[Default]")
+        reveal_type(x)  # N: Revealed type is "builtins.list[__main__.Custom]"
+        reveal_type(y)  # N: Revealed type is "builtins.list[builtins.int] | builtins.list[__main__.Default]"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.list[__main__.Custom] | __main__.Default"
+        reveal_type(y)  # N: Revealed type is "builtins.list[builtins.int] | builtins.list[__main__.Default]"
+
+def f3(x: Custom | dict[str, str], y: dict[int, int]):
+    if x == y:
+        reveal_type(x)  # N: Revealed type is "__main__.Custom | builtins.dict[builtins.str, builtins.str]"
+        reveal_type(y)  # N: Revealed type is "builtins.dict[builtins.int, builtins.int]"
+    else:
+        reveal_type(x)  # N: Revealed type is "__main__.Custom | builtins.dict[builtins.str, builtins.str]"
+        reveal_type(y)  # N: Revealed type is "builtins.dict[builtins.int, builtins.int]"
+[builtins fixtures/primitives.pyi]
+
+[case testNarrowingRecursiveCallable]
+# flags: --strict-equality --warn-unreachable
+from __future__ import annotations
+from typing import Callable
+
+class A: ...
+class B: ...
+
+T = Callable[[A], "S"]
+S = Callable[[B], "T"]
+
+def f(x: S, y: T):
+    if x == y:  # E: Unsupported left operand type for == ("Callable[[B], T]")
+        reveal_type(x)  # N: Revealed type is "def (__main__.B) -> def (__main__.A) -> ..."
+        reveal_type(y)  # N: Revealed type is "def (__main__.A) -> def (__main__.B) -> ..."
+    else:
+        reveal_type(x)  # N: Revealed type is "def (__main__.B) -> def (__main__.A) -> ..."
+        reveal_type(y)  # N: Revealed type is "def (__main__.A) -> def (__main__.B) -> ..."
+[builtins fixtures/tuple.pyi]
+
+[case testNarrowingRecursiveUnion]
+# flags: --strict-equality --warn-unreachable
+from __future__ import annotations
+from typing import Union
+
+class A: ...
+class B: ...
+
+T = Union[A, "S"]
+S = Union[B, "T"]  # E: Invalid recursive alias: a union item of itself
+
+def f(x: S, y: T):
+    if x == y:
+        reveal_type(x)  # N: Revealed type is "Any"
+        reveal_type(y)  # N: Revealed type is "__main__.A | Any"
+[builtins fixtures/tuple.pyi]
+
 [case testNarrowingUnreachableCases]
 # flags: --strict-equality --warn-unreachable
 from typing import Literal, Union