Put it together with other caches

Author: ilevkivskyi

mypy/typeops.py

@@ -9,7 +9,7 @@
 
 import itertools
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Final, TypeVar, cast
+from typing import Any, TypeVar, cast
 
 import mypy.subtypes
 from mypy.checker_state import checker_state
@@ -72,15 +72,10 @@
     instance_cache,
     remove_dups,
 )
+from mypy.typestate import BIND_SELF_CACHE_MAX_SIZE, type_state
 from mypy.typetraverser import TypeTraverserVisitor
 from mypy.typevars import fill_typevars
 
-BindSelfCacheKey = tuple[CallableType, Type, bool, bool]
-BindSelfCacheValue = tuple[CallableType, Sequence[TypeVarLikeType]]
-
-BIND_SELF_CACHE: Final[dict[BindSelfCacheKey, BindSelfCacheValue]] = {}
-BIND_SELF_CACHE_MAX_SIZE: Final = 10000
-
 
 def is_recursive_pair(s: Type, t: Type) -> bool:
     """Is this a pair of recursive types?
@@ -562,8 +557,8 @@ def bind_self_inner(
             # Only use cache for something that in not freshened (e.g. protocol access),
             # otherwise it would be pointless as each fresh meta variable is unique.
             cache_key = (func, original_type, is_classmethod, ignore_instances)
-        if cache_key is not None and cache_key in BIND_SELF_CACHE:
-            return BIND_SELF_CACHE[cache_key]
+        if cache_key is not None and cache_key in type_state.bind_self_cache:
+            return type_state.bind_self_cache[cache_key]
 
         # Solve for these type arguments using the actual class or instance type.
         typeargs = infer_type_arguments(
@@ -589,9 +584,9 @@ def bind_self_inner(
         func = expand_type(func, {tv.id: arg for tv, arg in zip(self_vars, to_apply)})
         variables = [v for v in func.variables if v not in self_vars]
         if cache_key is not None:
-            if len(BIND_SELF_CACHE) > BIND_SELF_CACHE_MAX_SIZE:
-                BIND_SELF_CACHE.clear()
-            BIND_SELF_CACHE[cache_key] = func, variables
+            if len(type_state.bind_self_cache) > BIND_SELF_CACHE_MAX_SIZE:
+                type_state.bind_self_cache.clear()
+            type_state.bind_self_cache[cache_key] = func, variables
     else:
         variables = func.variables
     return func, variables

mypy/typestate.py

@@ -5,15 +5,29 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from typing import Final, TypeAlias as _TypeAlias
 
 from mypy.nodes import VARIANCE_NOT_READY, TypeInfo
 from mypy.server.trigger import make_trigger
-from mypy.types import Instance, Type, TypeVarId, TypeVarType, get_proper_type
+from mypy.types import (
+    CallableType,
+    Instance,
+    Type,
+    TypeVarId,
+    TypeVarLikeType,
+    TypeVarType,
+    get_proper_type,
+)
 
 MAX_NEGATIVE_CACHE_TYPES: Final = 1000
 MAX_NEGATIVE_CACHE_ENTRIES: Final = 10000
 
+BindSelfCacheKey = tuple[CallableType, Type, bool, bool]
+BindSelfCacheValue = tuple[CallableType, Sequence[TypeVarLikeType]]
+
+BIND_SELF_CACHE_MAX_SIZE: Final = 10000
+
 # Represents that the 'left' instance is a subtype of the 'right' instance
 SubtypeRelationship: _TypeAlias = tuple[Instance, Instance]
 
@@ -47,6 +61,10 @@ class TypeState:
     # Same as above but for negative subtyping results.
     _negative_subtype_caches: Final[SubtypeCache]
 
+    # A dedicated lightweight cache for bind_self(), which may be slow because of type
+    # inference in some scenarios involving protocols with non-trivial self-types.
+    bind_self_cache: Final[dict[BindSelfCacheKey, BindSelfCacheValue]]
+
     # This contains protocol dependencies generated after running a full build,
     # or after an update. These dependencies are special because:
     #   * They are a global property of the program; i.e. some dependencies for imported
@@ -119,6 +137,7 @@ def __init__(self) -> None:
         self.infer_unions = False
         self.infer_polymorphic = False
         self.keep_unions = False
+        self.bind_self_cache = {}
 
     def is_assumed_subtype(self, left: Type, right: Type) -> bool:
         for l, r in reversed(self._assuming):
@@ -331,4 +350,5 @@ def reset_global_state() -> None:
     """
     type_state.reset_all_subtype_caches()
     type_state.reset_protocol_deps()
+    type_state.bind_self_cache.clear()
     TypeVarId.next_raw_id = 1