diff --git a/mypy/typeops.py b/mypy/typeops.py index 09f418b129ed5..3b96faeb2e5b4 100644 --- a/mypy/typeops.py +++ b/mypy/typeops.py @@ -5,7 +5,7 @@ since these may assume that MROs are ready. """ -from typing import cast, Optional, List, Sequence, Set, Iterable, TypeVar +from typing import cast, Optional, List, Sequence, Set, Iterable, TypeVar, Tuple from typing_extensions import Type as TypingType import sys @@ -344,35 +344,41 @@ def make_simplified_union(items: Sequence[Type], from mypy.subtypes import is_proper_subtype removed = set() # type: Set[int] - - # Avoid slow nested for loop for Union of Literal of strings (issue #9169) - if all((isinstance(item, LiteralType) and - item.fallback.type.fullname == 'builtins.str') - for item in items): - seen = set() # type: Set[str] - for index, item in enumerate(items): + seen_fast = set() # type: Set[Tuple[str, str]] + + # NB: having a separate fast path for Union of Literal and slow path for other things + # would arguably be cleaner, however it breaks down when simplifying the Union of two + # different enum types as try_expanding_enum_to_union works recursively and will + # trigger intermediate simplifications that would render the fast path useless + # Besides, the blended loop allows us to gracefully handle unions involving none/optional + # without any extra logic. + for i, item in enumerate(items): + if i in removed: + continue + # Avoid slow nested for loop for Union of Literal of strings/enums (issue #9169) + if isinstance(item, LiteralType) and ( + item.fallback.type.is_enum or item.fallback.type.fullname == 'builtins.str' + ): assert isinstance(item, LiteralType) assert isinstance(item.value, str) - if item.value in seen: - removed.add(index) - seen.add(item.value) - - else: - for i, ti in enumerate(items): - if i in removed: continue - # Keep track of the truishness info for deleted subtypes which can be relevant - cbt = cbf = False - for j, tj in enumerate(items): - if i != j and is_proper_subtype(tj, ti, keep_erased_types=keep_erased): - # We found a redundant item in the union. - removed.add(j) - cbt = cbt or tj.can_be_true - cbf = cbf or tj.can_be_false - # if deleted subtypes had more general truthiness, use that - if not ti.can_be_true and cbt: - items[i] = true_or_false(ti) - elif not ti.can_be_false and cbf: - items[i] = true_or_false(ti) + k = (item.value, item.fallback.type.fullname) + if k in seen_fast: + removed.add(i) + seen_fast.add(k) + continue + # Keep track of the truishness info for deleted subtypes which can be relevant + cbt = cbf = False + for j, tj in enumerate(items[i+1:], start=i+1): + if not isinstance(tj, LiteralType) and is_proper_subtype(tj, item, keep_erased_types=keep_erased): + # We found a redundant item in the union. + removed.add(j) + cbt = cbt or tj.can_be_true + cbf = cbf or tj.can_be_false + # if deleted subtypes had more general truthiness, use that + if not item.can_be_true and cbt: + items[i] = true_or_false(item) + elif not item.can_be_false and cbf: + items[i] = true_or_false(item) simplified_set = [items[i] for i in range(len(items)) if i not in removed] return UnionType.make_union(simplified_set, line, column)