Avoid unnecessary sum type expansion when narrowing (#22252)

conditional_types already expands bool and enum types when the target is
a matching literal.

Authored by Codex

Author: hauntsaninja

mypy/checker.py

@@ -7157,9 +7157,7 @@ def narrow_type_by_identity_equality(
                 if narrowable_expr_type is None:
                     if_type = else_type = ambiguous_expr_type
                 else:
-                    narrowable_expr_type = try_expanding_sum_type_to_union(
-                        coerce_to_literal(narrowable_expr_type), None
-                    )
+                    narrowable_expr_type = coerce_to_literal(narrowable_expr_type)
                     if_type, else_type = conditional_types(
                         narrowable_expr_type,
                         [TypeRange(target_type, is_upper_bound=False)],
@@ -7236,7 +7234,7 @@ def narrow_type_by_identity_equality(
                     # we narrow to in the if_map
                     or_if_maps.append({operands[i]: expr_type})
 
-                expr_type = coerce_to_literal(try_expanding_sum_type_to_union(expr_type, None))
+                expr_type = coerce_to_literal(expr_type)
                 for j in expr_indices:
                     if j in custom_eq_indices:
                         if i == j:
@@ -7262,9 +7260,6 @@ def narrow_type_by_identity_equality(
                     if narrowable_expr_type is None:
                         if_type = else_type = ambiguous_expr_type
                     else:
-                        narrowable_expr_type = coerce_to_literal(
-                            try_expanding_sum_type_to_union(narrowable_expr_type, None)
-                        )
                         if_type, else_type = conditional_types(
                             narrowable_expr_type,
                             [TypeRange(target_type, is_upper_bound=False)],

test-data/unit/check-narrowing.test

@@ -861,7 +861,7 @@ def f2(x: Union[Default, Literal[1], Literal[2]]):
 
 [case testNarrowingEqualityCustomEqualityEnum]
 # flags: --strict-equality --warn-unreachable
-from typing import Literal, Union
+from typing import Literal
 from enum import Enum
 
 class CustomEnum(Enum):
@@ -870,18 +870,24 @@ class CustomEnum(Enum):
 
     def __eq__(self, other: object) -> bool: return True
 
-x3: CustomEnum
-key: Literal[CustomEnum.A]
-if x3 == key:
-    reveal_type(x3)  # N: Revealed type is "__main__.CustomEnum"
-else:
-    reveal_type(x3)  # N: Revealed type is "Literal[__main__.CustomEnum.B]"
+def check_equality(x: CustomEnum, key: Literal[CustomEnum.A]) -> None:
+    if x == key:
+        reveal_type(x)  # N: Revealed type is "__main__.CustomEnum"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal[__main__.CustomEnum.B]"
 
 # For comparison, this narrows since we bypass __eq__
-if x3 is key:
-    reveal_type(x3)  # N: Revealed type is "Literal[__main__.CustomEnum.A]"
-else:
-    reveal_type(x3)  # N: Revealed type is "Literal[__main__.CustomEnum.B]"
+def check_identity(x: CustomEnum, key: Literal[CustomEnum.A]) -> None:
+    if x is key:
+        reveal_type(x)  # N: Revealed type is "Literal[__main__.CustomEnum.A]"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal[__main__.CustomEnum.B]"
+
+def check_optional_equality(x: CustomEnum | None, key: Literal[CustomEnum.A]) -> None:
+    if x == key:
+        reveal_type(x)  # N: Revealed type is "__main__.CustomEnum"
+    else:
+        reveal_type(x)  # N: Revealed type is "Literal[__main__.CustomEnum.B] | None"
 [builtins fixtures/primitives.pyi]
 
 [case testNarrowingEqualityCustomEqualityChainedComparison]