Speed up bind_self() in edge cases involving protocols

mypy/checkmember.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
-from typing import TypeVar, cast
+from typing import cast
 
 from mypy import message_registry, state
 from mypy.checker_shared import TypeCheckerSharedApi
@@ -21,7 +21,6 @@
 from mypy.nodes import (
     ARG_POS,
     ARG_STAR,
-    ARG_STAR2,
     EXCLUDED_ENUM_ATTRIBUTES,
     SYMBOL_FUNCBASE_TYPES,
     Context,
@@ -45,6 +44,7 @@
 from mypy.subtypes import is_subtype
 from mypy.typeops import (
     bind_self,
+    bind_self_fast,
     erase_to_bound,
     freeze_all_type_vars,
     get_all_type_vars,
@@ -1503,34 +1503,6 @@ def analyze_decorator_or_funcbase_access(
     return bind_self(typ, original_type=mx.self_type, is_classmethod=defn.is_class)
 
 
-F = TypeVar("F", bound=FunctionLike)
-
-
-def bind_self_fast(method: F, original_type: Type | None = None) -> F:
-    """Return a copy of `method`, with the type of its first parameter (usually
-    self or cls) bound to original_type.
-
-    This is a faster version of mypy.typeops.bind_self() that can be used for methods
-    with trivial self/cls annotations.
-    """
-    if isinstance(method, Overloaded):
-        items = [bind_self_fast(c, original_type) for c in method.items]
-        return cast(F, Overloaded(items))
-    assert isinstance(method, CallableType)
-    if not method.arg_types:
-        # Invalid method, return something.
-        return method
-    if method.arg_kinds[0] in (ARG_STAR, ARG_STAR2):
-        # See typeops.py for details.
-        return method
-    return method.copy_modified(
-        arg_types=method.arg_types[1:],
-        arg_kinds=method.arg_kinds[1:],
-        arg_names=method.arg_names[1:],
-        is_bound=True,
-    )
-
-
 def has_operator(typ: Type, op_method: str) -> bool:
     """Does type have operator with the given name?
 

mypy/server/update.py

@@ -887,6 +887,7 @@ def propagate_changes_using_dependencies(
         # unchanged, but was marked stale because its sub- (or super-) type changed.
         for info in stale_protos:
             type_state.reset_subtype_caches_for(info)
+        type_state.bind_self_cache.clear()
         # Then fully reprocess all targets.
         # TODO: Preserve order (set is not optimal)
         for id, nodes in sorted(todo.items(), key=lambda x: x[0]):

mypy/subtypes.py

@@ -30,6 +30,7 @@
     Context,
     Decorator,
     FuncBase,
+    FuncDef,
     OverloadedFuncDef,
     TypeInfo,
     Var,
@@ -1489,6 +1490,7 @@ def find_member(
         # to avoid type variable id collisions, but for protocols this means we can't
         # use the assumption stack, that will grow indefinitely.
         # TODO: find a cleaner solution that doesn't involve massive perf impact.
+        # Note: bind_self() cache relies on this, update it if this is changed.
         preserve_type_var_ids=True,
     )
     with type_checker.msg.filter_errors(filter_deprecated=True):
@@ -1645,8 +1647,6 @@ def find_node_type(
     """Find type of a variable or method 'node' (maybe also a decorated method).
     Apply type arguments from 'itype', and bind 'self' to 'subtype'.
     """
-    from mypy.typeops import bind_self
-
     if isinstance(node, FuncBase):
         typ: Type | None = mypy.typeops.function_type(
             node, fallback=Instance(itype.type.mro[-1], [])
@@ -1676,8 +1676,10 @@ def find_node_type(
         ):
             # Don't bind instance methods on class objects.
             signature = p_typ
+        elif isinstance(node, (FuncDef, OverloadedFuncDef)) and node.is_trivial_self:
+            signature = mypy.typeops.bind_self_fast(p_typ, subtype)
         else:
-            signature = bind_self(
+            signature = mypy.typeops.bind_self(
                 p_typ, subtype, is_classmethod=isinstance(node, Var) and node.is_classmethod
             )
         if node.is_property and not class_obj:

mypy/typeops.py

@@ -72,6 +72,7 @@
     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
 
@@ -551,6 +552,14 @@ def bind_self_inner(
         self_ids = {tv.id for tv in get_all_type_vars(self_param_type)}
         self_vars = [tv for tv in func.variables if tv.id in self_ids]
 
+        cache_key = None
+        if not any(tv.id.is_meta_var() for tv in self_vars) and state.strict_optional:
+            # 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 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(
             self_vars, self_param_type, original_type, is_supertype=True, erase_types=False
@@ -574,11 +583,40 @@ def bind_self_inner(
         to_apply = [t if t is not None else UninhabitedType() for t in typeargs]
         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(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
 
 
+def bind_self_fast(method: F, original_type: Type | None = None) -> F:
+    """Return a copy of `method`, with the type of its first parameter (usually
+    self or cls) bound to original_type.
+
+    This is a faster version of mypy.typeops.bind_self() that can be used for methods
+    with trivial self/cls annotations.
+    """
+    if isinstance(method, Overloaded):
+        items = [bind_self_fast(c, original_type) for c in method.items]
+        return cast(F, Overloaded(items))
+    assert isinstance(method, CallableType)
+    if not method.arg_types:
+        # Invalid method, return something.
+        return method
+    if method.arg_kinds[0] in (ARG_STAR, ARG_STAR2):
+        # See bind_self() for details.
+        return method
+    return method.copy_modified(
+        arg_types=method.arg_types[1:],
+        arg_kinds=method.arg_kinds[1:],
+        arg_names=method.arg_names[1:],
+        is_bound=True,
+    )
+
+
 def erase_to_bound(t: Type) -> Type:
     # TODO: use value restrictions to produce a union?
     t = get_proper_type(t)

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

test-data/unit/fine-grained.test

@@ -11970,3 +11970,39 @@ class C: pass
 [out]
 ==
 main:3: error: Name "C" is not defined
+
+[case testProtocolGenericSelfTypeBaseChanged]
+import a
+[file a.py]
+from b import P
+from d import C
+def f(x: P) -> None: pass
+f(C())
+[file b.py]
+from typing import Protocol
+class P(Protocol):
+    def meth(self) -> int: ...
+[file c.py]
+from typing import Generic, TypeVar
+T = TypeVar("T")
+S = TypeVar("S")
+class Base(Generic[T]):
+    def meth(self: "Base[S]") -> S: ...
+[file d.py]
+from c import Base
+class C(Base[int]): pass
+[file d.py.2]
+from c import Base
+class C(Base[str]): pass
+[file d.py.3]
+from c import Base
+class C(Base[int]): pass
+[out]
+==
+a.py:4: error: Argument 1 to "f" has incompatible type "C"; expected "P"
+a.py:4: note: Following member(s) of "C" have conflicts:
+a.py:4: note:     Expected:
+a.py:4: note:         def meth(self) -> int
+a.py:4: note:     Got:
+a.py:4: note:         def meth(self) -> str
+==