From 6d2f05a0c3a3d669f92633ad8270da500b56b3a8 Mon Sep 17 00:00:00 2001 From: Shantanu Jain Date: Wed, 21 Oct 2026 06:30:36 -0700 Subject: [PATCH] Avoid unnecessary sum type expansion when narrowing conditional_types already expands bool and enum types when the target is a matching literal. Authored by Codex --- mypy/checker.py | 9 ++------- test-data/unit/check-narrowing.test | 28 +++++++++++++++++----------- 2 files changed, 19 insertions(+), 18 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 64852cb724f4..2d29f6aa43ae 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -7168,9 +7168,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)], @@ -7247,7 +7245,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: @@ -7273,9 +7271,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)], diff --git a/test-data/unit/check-narrowing.test b/test-data/unit/check-narrowing.test index 1ee146be86f0..38ff124518d2 100644 --- a/test-data/unit/check-narrowing.test +++ b/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]