Skip to content

Commit

Permalink
typeops: extend make_simplified_union fast path to enums
Browse files Browse the repository at this point in the history
In PR python#9192 a fast path was created to address the slowness reported
in issue python#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.
  • Loading branch information
huguesb committed Sep 1, 2020
1 parent 13ae58f commit c08e7be
Showing 1 changed file with 22 additions and 4 deletions.
26 changes: 22 additions & 4 deletions mypy/typeops.py
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit c08e7be

Please sign in to comment.