Speed up freshening type variables

mypy/expandtype.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 from typing import Iterable, Mapping, Sequence, TypeVar, cast, overload
+from typing_extensions import Final
 
 from mypy.nodes import ARG_POS, ARG_STAR, Var
 from mypy.type_visitor import TypeTranslator
 from mypy.types import (
+    ANY_STRATEGY,
     AnyType,
+    BoolTypeQuery,
     CallableType,
     DeletedType,
     ErasedType,
@@ -138,13 +141,30 @@ def freshen_function_type_vars(callee: F) -> F:
         return cast(F, fresh_overload)
 
 
+class HasGenericCallable(BoolTypeQuery):
+    def __init__(self) -> None:
+        super().__init__(ANY_STRATEGY)
+
+    def visit_callable_type(self, t: CallableType) -> bool:
+        return t.is_generic() or super().visit_callable_type(t)
+
+
+# Share a singleton since this is performance sensitive
+has_generic_callable: Final = HasGenericCallable()
+
+
 T = TypeVar("T", bound=Type)
 
 
 def freshen_all_functions_type_vars(t: T) -> T:
-    result = t.accept(FreshenCallableVisitor())
-    assert isinstance(result, type(t))
-    return result
+    result: Type
+    has_generic_callable.reset()
+    if not t.accept(has_generic_callable):
+        return t  # Fast path to avoid expensive freshening
+    else:
+        result = t.accept(FreshenCallableVisitor())
+        assert isinstance(result, type(t))
+        return result
 
 
 class FreshenCallableVisitor(TypeTranslator):

mypy/type_visitor.py

@@ -15,6 +15,7 @@
 
 from abc import abstractmethod
 from typing import Any, Callable, Generic, Iterable, Sequence, TypeVar, cast
+from typing_extensions import Final
 
 from mypy_extensions import mypyc_attr, trait
 
@@ -417,3 +418,158 @@ def visit_type_alias_type(self, t: TypeAliasType) -> T:
     def query_types(self, types: Iterable[Type]) -> T:
         """Perform a query for a list of types using the strategy to combine the results."""
         return self.strategy([t.accept(self) for t in types])
+
+
+# Return True if at least one type component returns True
+ANY_STRATEGY: Final = 0
+# Return True if no type component returns False
+ALL_STRATEGY: Final = 1
+
+
+class BoolTypeQuery(SyntheticTypeVisitor[bool]):
+    """Visitor for performing recursive queries of types with a bool result.
+
+    Use TypeQuery if you need non-bool results.
+
+    'strategy' is used to combine results for a series of types. It must
+    be ANY_STRATEGY or ALL_STRATEGY.
+
+    Note: This visitor keeps an internal state (tracks type aliases to avoid
+    recursion), so it should *never* be re-used for querying different types
+    unless you call reset() first.
+    """
+
+    def __init__(self, strategy: int) -> None:
+        self.strategy = strategy
+        if strategy == ANY_STRATEGY:
+            self.default = False
+        else:
+            assert strategy == ALL_STRATEGY
+            self.default = True
+        # Keep track of the type aliases already visited. This is needed to avoid
+        # infinite recursion on types like A = Union[int, List[A]]. An empty set is
+        # represented as None as a micro-optimization.
+        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
+
+    def reset(self) -> None:
+        """Clear mutable state (but preserve strategy).
+
+        This *must* be called if you want to reuse the visitor.
+        """
+        self.seen_aliases = None
+
+    def visit_unbound_type(self, t: UnboundType) -> bool:
+        return self.query_types(t.args)
+
+    def visit_type_list(self, t: TypeList) -> bool:
+        return self.query_types(t.items)
+
+    def visit_callable_argument(self, t: CallableArgument) -> bool:
+        return t.typ.accept(self)
+
+    def visit_any(self, t: AnyType) -> bool:
+        return self.default
+
+    def visit_uninhabited_type(self, t: UninhabitedType) -> bool:
+        return self.default
+
+    def visit_none_type(self, t: NoneType) -> bool:
+        return self.default
+
+    def visit_erased_type(self, t: ErasedType) -> bool:
+        return self.default
+
+    def visit_deleted_type(self, t: DeletedType) -> bool:
+        return self.default
+
+    def visit_type_var(self, t: TypeVarType) -> bool:
+        return self.query_types([t.upper_bound] + t.values)
+
+    def visit_param_spec(self, t: ParamSpecType) -> bool:
+        return self.default
+
+    def visit_type_var_tuple(self, t: TypeVarTupleType) -> bool:
+        return self.default
+
+    def visit_unpack_type(self, t: UnpackType) -> bool:
+        return self.query_types([t.type])
+
+    def visit_parameters(self, t: Parameters) -> bool:
+        return self.query_types(t.arg_types)
+
+    def visit_partial_type(self, t: PartialType) -> bool:
+        return self.default
+
+    def visit_instance(self, t: Instance) -> bool:
+        return self.query_types(t.args)
+
+    def visit_callable_type(self, t: CallableType) -> bool:
+        # FIX generics
+        # Avoid allocating any objects here as an optimization.
+        args = self.query_types(t.arg_types)
+        ret = t.ret_type.accept(self)
+        if self.strategy == ANY_STRATEGY:
+            return args or ret
+        else:
+            return args and ret
+
+    def visit_tuple_type(self, t: TupleType) -> bool:
+        return self.query_types(t.items)
+
+    def visit_typeddict_type(self, t: TypedDictType) -> bool:
+        return self.query_types(list(t.items.values()))
+
+    def visit_raw_expression_type(self, t: RawExpressionType) -> bool:
+        return self.default
+
+    def visit_literal_type(self, t: LiteralType) -> bool:
+        return self.default
+
+    def visit_star_type(self, t: StarType) -> bool:
+        return t.type.accept(self)
+
+    def visit_union_type(self, t: UnionType) -> bool:
+        return self.query_types(t.items)
+
+    def visit_overloaded(self, t: Overloaded) -> bool:
+        return self.query_types(t.items)  # type: ignore[arg-type]
+
+    def visit_type_type(self, t: TypeType) -> bool:
+        return t.item.accept(self)
+
+    def visit_ellipsis_type(self, t: EllipsisType) -> bool:
+        return self.default
+
+    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)
+        return get_proper_type(t).accept(self)
+
+    def query_types(self, types: list[Type] | tuple[Type, ...]) -> bool:
+        """Perform a query for a sequence of types using the strategy to combine the results."""
+        # Special-case for lists and tuples to allow mypyc to produce better code.
+        if isinstance(types, list):
+            if self.strategy == ANY_STRATEGY:
+                return any(t.accept(self) for t in types)
+            else:
+                return all(t.accept(self) for t in types)
+        else:
+            if self.strategy == ANY_STRATEGY:
+                return any(t.accept(self) for t in types)
+            else:
+                return all(t.accept(self) for t in types)

mypy/types.py

@@ -2874,7 +2874,10 @@ def get_proper_types(it: Iterable[Type | None]) -> list[ProperType] | list[Prope
 # to make it easier to gradually get modules working with mypyc.
 # Import them here, after the types are defined.
 # This is intended as a re-export also.
-from mypy.type_visitor import (  # noqa: F811
+from mypy.type_visitor import (  # noqa: F811,F401
+    ALL_STRATEGY as ALL_STRATEGY,
+    ANY_STRATEGY as ANY_STRATEGY,
+    BoolTypeQuery as BoolTypeQuery,
     SyntheticTypeVisitor as SyntheticTypeVisitor,
     TypeQuery as TypeQuery,
     TypeTranslator as TypeTranslator,