diff --git a/mypy/meet.py b/mypy/meet.py index bfc6d88a1e20..05190daffba7 100644 --- a/mypy/meet.py +++ b/mypy/meet.py @@ -114,6 +114,66 @@ def meet_types(s: Type, t: Type) -> ProperType: return t.accept(TypeMeetVisitor(s)) +def _is_simple_union_item(t: ProperType) -> bool: + """Items that overlap another simple item iff they compare equal.""" + return isinstance(t, (LiteralType, NoneType, UninhabitedType)) + + +def _union_items_may_narrow(declared: Type, narrowed: Type) -> bool: + """Keep the historical union-narrowing overlap predicate. + + This special-casing is needed to support checking branches like this: + + x: Union[float, complex] + if isinstance(x, int): + ... + + And assignments like this: + + x: float | None + y: int | None + x = y + """ + return is_overlapping_types(declared, narrowed, ignore_promotions=True) or is_subtype( + narrowed, declared, ignore_promotions=False + ) + + +def _narrow_declared_union_items( + declared_items: list[Type], narrowed_items: list[Type] +) -> ProperType: + """Intersect two unions, hashing simple items to avoid O(N*M) overlap checks. + + After equality narrowing, enums explode into unions of literals. Intersecting + two such unions pairwise is quadratic and dominates match-statement checking + on large enums (see #21997). Literal/None/Never items can be intersected via + hashing; remaining items still use the original overlap predicate. + """ + simple_narrowed: dict[ProperType, Type] = {} + complex_narrowed: list[Type] = [] + for item in narrowed_items: + proper = get_proper_type(item) + if _is_simple_union_item(proper): + simple_narrowed[proper] = item + else: + complex_narrowed.append(item) + + result: list[Type] = [] + for declared in declared_items: + proper = get_proper_type(declared) + if _is_simple_union_item(proper): + if proper in simple_narrowed: + result.append(declared) + for narrowed in complex_narrowed: + if _union_items_may_narrow(declared, narrowed): + result.append(narrow_declared_type(declared, narrowed)) + else: + for narrowed in narrowed_items: + if _union_items_may_narrow(declared, narrowed): + result.append(narrow_declared_type(declared, narrowed)) + return make_simplified_union(result) + + def narrow_declared_type(declared: Type, narrowed: Type) -> Type: """Return the declared type narrowed down to another type.""" # TODO: check infinite recursion for aliases here. @@ -133,30 +193,12 @@ def narrow_declared_type(declared: Type, narrowed: Type) -> Type: return original_declared if isinstance(declared, UnionType): declared_items = declared.relevant_items() + narrowed_items: list[Type] if isinstance(narrowed, UnionType): narrowed_items = narrowed.relevant_items() else: narrowed_items = [narrowed] - return make_simplified_union( - [ - narrow_declared_type(d, n) - for d in declared_items - for n in narrowed_items - # This (ugly) special-casing is needed to support checking - # branches like this: - # x: Union[float, complex] - # if isinstance(x, int): - # ... - # And assignments like this: - # x: float | None - # y: int | None - # x = y - if ( - is_overlapping_types(d, n, ignore_promotions=True) - or is_subtype(n, d, ignore_promotions=False) - ) - ] - ) + return _narrow_declared_union_items(declared_items, narrowed_items) if is_enum_overlapping_union(declared, narrowed): # Quick check before reaching `is_overlapping_types`. If it's enum/literal overlap, # avoid full expansion and make it faster. diff --git a/mypy/test/testtypes.py b/mypy/test/testtypes.py index b287e82b3d4a..598cf8cd674e 100644 --- a/mypy/test/testtypes.py +++ b/mypy/test/testtypes.py @@ -1359,6 +1359,7 @@ def test_type_type(self) -> None: def test_literal_type(self) -> None: a = self.fx.a + b = self.fx.b lit1 = self.fx.lit1 lit2 = self.fx.lit2 lit3 = self.fx.lit3 @@ -1376,6 +1377,30 @@ def test_literal_type(self) -> None: assert is_same_type(lit1, narrow_declared_type(lit1, a)) assert is_same_type(lit2, narrow_declared_type(lit2, a)) + # Union-union narrowing keeps overlapping literals and drops the rest. + assert is_same_type( + lit2, narrow_declared_type(UnionType([lit1, lit2]), UnionType([lit2, lit3])) + ) + # Mixed literal/instance unions still use the overlap predicate. + assert is_same_type( + UnionType([lit1, b]), narrow_declared_type(UnionType([lit1, a]), UnionType([lit1, b])) + ) + + def test_narrow_declared_type_large_literal_unions(self) -> None: + # Equality narrowing explodes enums into large literal unions. Intersecting + # those unions used to be quadratic; this should stay effectively linear. + fx = self.fx + left = [LiteralType("v%d" % i, fx.str_type) for i in range(2000)] + right = [LiteralType("v%d" % i, fx.str_type) for i in range(500, 2500)] + result = get_proper_type(narrow_declared_type(UnionType(left), UnionType(right))) + assert isinstance(result, UnionType) + values = set() + for item in result.items: + proper = get_proper_type(item) + if isinstance(proper, LiteralType): + values.add(proper.value) + assert values == {"v%d" % i for i in range(500, 2000)} + # FIX generic interfaces + ranges def assert_meet_uninhabited(self, s: Type, t: Type) -> None: diff --git a/test-data/unit/check-python310.test b/test-data/unit/check-python310.test index 01e490d0da50..3c5e36defa8d 100644 --- a/test-data/unit/check-python310.test +++ b/test-data/unit/check-python310.test @@ -2310,6 +2310,25 @@ def g(m: Medal) -> int: return 2 [builtins fixtures/enum.pyi] +[case testMatchAfterEqualityNarrowingEnum] +# flags: --strict-equality --warn-unreachable +from enum import Enum + +class Medal(Enum): + gold = 1 + silver = 2 + bronze = 3 + +def f(m: Medal) -> None: + if m == Medal.gold: + return + match m: + case Medal.silver: + reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.silver]" + case Medal.bronze: + reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.bronze]" +[builtins fixtures/primitives.pyi] + [case testMatchLiteralOrValuePattern] # flags: --strict-equality --warn-unreachable from typing import Literal