From c08e7be23184f503985d037b1df9bafb9bfbdae2 Mon Sep 17 00:00:00 2001 From: Hugues Bruant Date: Tue, 1 Sep 2020 16:56:08 -0700 Subject: [PATCH] typeops: extend make_simplified_union fast path to enums In PR #9192 a fast path was created to address the slowness reported in issue #9169 wherein large Union or literal types would dramatically slow down typechecking. It is desirable to extend this fast path to cover Enum types, as these can also leverage the O(n) set-based fast path instead of the O(n**2) fallback. This is seen to bring down the typechecking of a single fairly simple chain of `if` statements operating on a large enum (~3k members) from ~40min to 12s in real-world code! Note that the timing is taken from a pure-python run of mypy, as opposed to a compiled version. --- mypy/typeops.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/mypy/typeops.py b/mypy/typeops.py index 09f418b129ed5..a973b2cb2165c 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, Union, TypeVar from typing_extensions import Type as TypingType import sys @@ -311,6 +311,26 @@ def callable_corresponding_argument(typ: CallableType, return by_name if by_name is not None else by_pos +def _can_use_fast_path(items: Union[List[ProperType], List[Optional[ProperType]]]) -> bool: + types = set() # type: Set[str] + + for item in items: + if isinstance(item, LiteralType): + item = item.fallback + if not isinstance(item, Instance): + return False + if item.type.is_enum or item.type.fullname == 'builtins.str': + types.add(item.type.fullname) + else: + return False + + # we allow a fast path if and only if + # - all elements are of the same type + # - the singular type of all elements is an Enum or a string + # TODO: supports int as well? + return len(types) == 1 + + def make_simplified_union(items: Sequence[Type], line: int = -1, column: int = -1, *, keep_erased: bool = False) -> ProperType: @@ -346,9 +366,7 @@ def make_simplified_union(items: Sequence[Type], 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): + if _can_use_fast_path(items): seen = set() # type: Set[str] for index, item in enumerate(items): assert isinstance(item, LiteralType)