Try some aliases speed-up

Author: ilevkivskyi

mypy/constraints.py

@@ -21,6 +21,7 @@
     ArgKind,
     TypeInfo,
 )
+from mypy.type_visitor import ALL_STRATEGY, BoolTypeQuery
 from mypy.types import (
     TUPLE_LIKE_INSTANCE_NAMES,
     AnyType,
@@ -41,7 +42,6 @@
     TypeAliasType,
     TypedDictType,
     TypeOfAny,
-    TypeQuery,
     TypeType,
     TypeVarId,
     TypeVarLikeType,
@@ -670,9 +670,9 @@ def is_complete_type(typ: Type) -> bool:
     return typ.accept(CompleteTypeVisitor())
 
 
-class CompleteTypeVisitor(TypeQuery[bool]):
+class CompleteTypeVisitor(BoolTypeQuery):
     def __init__(self) -> None:
-        super().__init__(all)
+        super().__init__(ALL_STRATEGY)
 
     def visit_uninhabited_type(self, t: UninhabitedType) -> bool:
         return False

mypy/indirection.py

@@ -37,16 +37,15 @@ def find_modules(self, typs: Iterable[types.Type]) -> set[str]:
         return self.modules
 
     def _visit(self, typ: types.Type) -> None:
-        if isinstance(typ, types.TypeAliasType):
+        if isinstance(typ, types.TypeAliasType) and typ.is_recursive:
             # Avoid infinite recursion for recursive type aliases.
-            if typ not in self.seen_aliases:
-                self.seen_aliases.add(typ)
+            self.seen_aliases.add(typ)
         typ.accept(self)
 
     def _visit_type_tuple(self, typs: tuple[types.Type, ...]) -> None:
         # Micro-optimization: Specialized version of _visit for lists
         for typ in typs:
-            if isinstance(typ, types.TypeAliasType):
+            if isinstance(typ, types.TypeAliasType) and typ.is_recursive:
                 # Avoid infinite recursion for recursive type aliases.
                 if typ in self.seen_aliases:
                     continue
@@ -56,7 +55,7 @@ def _visit_type_tuple(self, typs: tuple[types.Type, ...]) -> None:
     def _visit_type_list(self, typs: list[types.Type]) -> None:
         # Micro-optimization: Specialized version of _visit for tuples
         for typ in typs:
-            if isinstance(typ, types.TypeAliasType):
+            if isinstance(typ, types.TypeAliasType) and typ.is_recursive:
                 # Avoid infinite recursion for recursive type aliases.
                 if typ in self.seen_aliases:
                     continue

mypy/semanal_typeargs.py

@@ -83,12 +83,11 @@ def visit_block(self, o: Block) -> None:
 
     def visit_type_alias_type(self, t: TypeAliasType) -> None:
         super().visit_type_alias_type(t)
-        if t in self.seen_aliases:
-            # Avoid infinite recursion on recursive type aliases.
-            # Note: it is fine to skip the aliases we have already seen in non-recursive
-            # types, since errors there have already been reported.
-            return
-        self.seen_aliases.add(t)
+        if t.is_recursive:
+            if t in self.seen_aliases:
+                # Avoid infinite recursion on recursive type aliases.
+                return
+            self.seen_aliases.add(t)
         assert t.alias is not None, f"Unfixed type alias {t.type_ref}"
         is_error, is_invalid = self.validate_args(
             t.alias.name, tuple(t.args), t.alias.alias_tvars, t

mypy/stats.py

@@ -43,6 +43,7 @@
     YieldFromExpr,
 )
 from mypy.traverser import TraverserVisitor
+from mypy.type_visitor import ANY_STRATEGY, BoolTypeQuery
 from mypy.typeanal import collect_all_inner_types
 from mypy.types import (
     AnyType,
@@ -52,7 +53,6 @@
     TupleType,
     Type,
     TypeOfAny,
-    TypeQuery,
     TypeVarType,
     get_proper_type,
     get_proper_types,
@@ -453,9 +453,9 @@ def is_imprecise(t: Type) -> bool:
     return t.accept(HasAnyQuery())
 
 
-class HasAnyQuery(TypeQuery[bool]):
+class HasAnyQuery(BoolTypeQuery):
     def __init__(self) -> None:
-        super().__init__(any)
+        super().__init__(ANY_STRATEGY)
 
     def visit_any(self, t: AnyType) -> bool:
         return not is_special_form_any(t)

mypy/test/testtypes.py

@@ -201,18 +201,6 @@ def test_type_alias_expand_once(self) -> None:
         assert get_proper_type(A) == target
         assert get_proper_type(target) == target
 
-    def test_type_alias_expand_all(self) -> None:
-        A, _ = self.fx.def_alias_1(self.fx.a)
-        assert A.expand_all_if_possible() is None
-        A, _ = self.fx.def_alias_2(self.fx.a)
-        assert A.expand_all_if_possible() is None
-
-        B = self.fx.non_rec_alias(self.fx.a)
-        C = self.fx.non_rec_alias(TupleType([B, B], Instance(self.fx.std_tuplei, [B])))
-        assert C.expand_all_if_possible() == TupleType(
-            [self.fx.a, self.fx.a], Instance(self.fx.std_tuplei, [self.fx.a])
-        )
-
     def test_recursive_nested_in_non_recursive(self) -> None:
         A, _ = self.fx.def_alias_1(self.fx.a)
         T = TypeVarType(

mypy/type_visitor.py

@@ -15,7 +15,7 @@
 
 from abc import abstractmethod
 from collections.abc import Iterable, Sequence
-from typing import Any, Callable, Final, Generic, TypeVar, cast
+from typing import Any, Final, Generic, TypeVar, cast
 
 from mypy_extensions import mypyc_attr, trait
 
@@ -353,16 +353,19 @@ class TypeQuery(SyntheticTypeVisitor[T]):
     # TODO: check that we don't have existing violations of this rule.
     """
 
-    def __init__(self, strategy: Callable[[list[T]], T]) -> None:
-        self.strategy = strategy
+    def __init__(self) -> None:
         # Keep track of the type aliases already visited. This is needed to avoid
         # infinite recursion on types like A = Union[int, List[A]].
-        self.seen_aliases: set[TypeAliasType] = set()
+        self.seen_aliases: set[TypeAliasType] | None = None
         # By default, we eagerly expand type aliases, and query also types in the
         # alias target. In most cases this is a desired behavior, but we may want
         # to skip targets in some cases (e.g. when collecting type variables).
         self.skip_alias_target = False
 
+    @abstractmethod
+    def strategy(self, items: list[T]) -> T:
+        raise NotImplementedError
+
     def visit_unbound_type(self, t: UnboundType, /) -> T:
         return self.query_types(t.args)
 
@@ -440,14 +443,15 @@ def visit_placeholder_type(self, t: PlaceholderType, /) -> T:
         return self.query_types(t.args)
 
     def visit_type_alias_type(self, t: TypeAliasType, /) -> T:
-        # Skip type aliases already visited types to avoid infinite recursion.
-        # TODO: Ideally we should fire subvisitors here (or use caching) if we care
-        #       about duplicates.
-        if t in self.seen_aliases:
-            return self.strategy([])
-        self.seen_aliases.add(t)
         if self.skip_alias_target:
             return self.query_types(t.args)
+        # Skip type aliases already visited types to avoid infinite recursion.
+        if t.is_recursive:
+            if self.seen_aliases is None:
+                self.seen_aliases = set()
+            elif t in self.seen_aliases:
+                return self.strategy([])
+            self.seen_aliases.add(t)
         return get_proper_type(t).accept(self)
 
     def query_types(self, types: Iterable[Type]) -> T:
@@ -580,16 +584,15 @@ def visit_placeholder_type(self, t: PlaceholderType, /) -> bool:
         return self.query_types(t.args)
 
     def visit_type_alias_type(self, t: TypeAliasType, /) -> bool:
-        # Skip type aliases already visited types to avoid infinite recursion.
-        # TODO: Ideally we should fire subvisitors here (or use caching) if we care
-        #       about duplicates.
-        if self.seen_aliases is None:
-            self.seen_aliases = set()
-        elif t in self.seen_aliases:
-            return self.default
-        self.seen_aliases.add(t)
         if self.skip_alias_target:
             return self.query_types(t.args)
+        # Skip type aliases already visited types to avoid infinite recursion.
+        if t.is_recursive:
+            if self.seen_aliases is None:
+                self.seen_aliases = set()
+            elif t in self.seen_aliases:
+                return self.default
+            self.seen_aliases.add(t)
         return get_proper_type(t).accept(self)
 
     def query_types(self, types: list[Type] | tuple[Type, ...]) -> bool:

mypy/typeanal.py

@@ -2377,9 +2377,9 @@ def has_explicit_any(t: Type) -> bool:
     return t.accept(HasExplicitAny())
 
 
-class HasExplicitAny(TypeQuery[bool]):
+class HasExplicitAny(BoolTypeQuery):
     def __init__(self) -> None:
-        super().__init__(any)
+        super().__init__(ANY_STRATEGY)
 
     def visit_any(self, t: AnyType) -> bool:
         return t.type_of_any == TypeOfAny.explicit
@@ -2418,15 +2418,11 @@ def collect_all_inner_types(t: Type) -> list[Type]:
 
 
 class CollectAllInnerTypesQuery(TypeQuery[list[Type]]):
-    def __init__(self) -> None:
-        super().__init__(self.combine_lists_strategy)
-
     def query_types(self, types: Iterable[Type]) -> list[Type]:
         return self.strategy([t.accept(self) for t in types]) + list(types)
 
-    @classmethod
-    def combine_lists_strategy(cls, it: Iterable[list[Type]]) -> list[Type]:
-        return list(itertools.chain.from_iterable(it))
+    def strategy(self, items: Iterable[list[Type]]) -> list[Type]:
+        return list(itertools.chain.from_iterable(items))
 
 
 def make_optional_type(t: Type) -> Type:
@@ -2556,7 +2552,6 @@ def __init__(self, api: SemanticAnalyzerCoreInterface, scope: TypeVarLikeScope)
         self.scope = scope
         self.type_var_likes: list[tuple[str, TypeVarLikeExpr]] = []
         self.has_self_type = False
-        self.seen_aliases: set[TypeAliasType] | None = None
         self.include_callables = True
 
     def _seems_like_callable(self, type: UnboundType) -> bool:
@@ -2653,7 +2648,8 @@ def visit_union_type(self, t: UnionType) -> None:
         self.process_types(t.items)
 
     def visit_overloaded(self, t: Overloaded) -> None:
-        self.process_types(t.items)  # type: ignore[arg-type]
+        for it in t.items:
+            it.accept(self)
 
     def visit_type_type(self, t: TypeType) -> None:
         t.item.accept(self)
@@ -2665,12 +2661,6 @@ def visit_placeholder_type(self, t: PlaceholderType) -> None:
         return self.process_types(t.args)
 
     def visit_type_alias_type(self, t: TypeAliasType) -> None:
-        # Skip type aliases in already visited types to avoid infinite recursion.
-        if self.seen_aliases is None:
-            self.seen_aliases = set()
-        elif t in self.seen_aliases:
-            return
-        self.seen_aliases.add(t)
         self.process_types(t.args)
 
     def process_types(self, types: list[Type] | tuple[Type, ...]) -> None:

mypy/typeops.py

@@ -1114,12 +1114,12 @@ def get_all_type_vars(tp: Type) -> list[TypeVarLikeType]:
 
 class TypeVarExtractor(TypeQuery[list[TypeVarLikeType]]):
     def __init__(self, include_all: bool = False) -> None:
-        super().__init__(self._merge)
+        super().__init__()
         self.include_all = include_all
 
-    def _merge(self, iter: Iterable[list[TypeVarLikeType]]) -> list[TypeVarLikeType]:
+    def strategy(self, items: Iterable[list[TypeVarLikeType]]) -> list[TypeVarLikeType]:
         out = []
-        for item in iter:
+        for item in items:
             out.extend(item)
         return out