Try some aliases speed-up (#19810)

This makes self-check almost 2% faster on my desktop (Python 3.12,
compiled -O2). Inspired by a slight regression in
#19798 I decided to re-think how we
detect/label the recursive types.

The new algorithm is not 100% equivalent to old one, but should be much
faster. The main semantic difference is this:
```python
A = list[B1]
B1 = list[B2]
B2 = list[B1]
```
previously all three aliases where labeled as recursive, now only last
two are. Which is kind of correct if you think about it for some time,
there is nothing genuinely recursive in `A` by itself. As a result:
* We have somewhat more verbose `reveal_type()` for recursive types
after fine-grained increments. Excessive use of `get_proper_type()` in
the daemon code is a known issue. I will take a look at it when I will
have a chance.
* I cleaned up some of relevant visitors to be more consistent with
recursive aliases.
* I also do couple cleanups/speedups in the type queries while I am at
it.

If there are no comments/objections, I will merge it later today. Then I
will merge #19798, and then _maybe_
an equivalent optimization for recursive instances like `class
str(Sequence[str]): ...`

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

@@ -39,8 +39,7 @@ def find_modules(self, typs: Iterable[types.Type]) -> set[str]:
     def _visit(self, typ: types.Type) -> None:
         if isinstance(typ, types.TypeAliasType):
             # 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:

mypy/mixedtraverser.py

@@ -47,6 +47,8 @@ def visit_class_def(self, o: ClassDef, /) -> None:
         if info:
             for base in info.bases:
                 base.accept(self)
+            if info.special_alias:
+                info.special_alias.accept(self)
 
     def visit_type_alias_expr(self, o: TypeAliasExpr, /) -> None:
         super().visit_type_alias_expr(o)

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
@@ -101,9 +100,12 @@ def visit_type_alias_type(self, t: TypeAliasType) -> None:
         if not is_error:
             # If there was already an error for the alias itself, there is no point in checking
             # the expansion, most likely it will result in the same kind of error.
-            get_proper_type(t).accept(self)
-            if t.alias is not None:
-                t.alias.accept(self)
+            if t.args:
+                # Since we always allow unbounded type variables in alias definitions, we need
+                # to verify the arguments satisfy the upper bounds of the expansion as well.
+                get_proper_type(t).accept(self)
+        if t.is_recursive:
+            self.seen_aliases.discard(t)
 
     def visit_tuple_type(self, t: TupleType) -> None:
         t.items = flatten_nested_tuples(t.items)

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
+        # (also use this as a simple-minded cache).
+        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.skip_alias_target:
+            return self.query_types(t.args)
+        # Skip type aliases already visited types to avoid infinite recursion
+        # (also use this as a simple-minded cache).
         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)
         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