Speed up bind_self() in trivial cases (#19024)

See #18991 for context.

We can skip all of logic in `check_self_arg()` and 90% of logic in
`bind_self()` for methods with trivial `self`/`cls` (i.e. if first
argument has no explicit annotation and there is no `Self` in
signature).

Locally I see 3-4% performance improvement (for self-check with
non-compiled mypy).

Author: ilevkivskyi

mypy/checkmember.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Callable, cast
+from typing import Callable, TypeVar, cast
 
 from mypy import message_registry, state, subtypes
 from mypy.checker_shared import TypeCheckerSharedApi
@@ -18,6 +18,7 @@
 from mypy.nodes import (
     ARG_POS,
     ARG_STAR,
+    ARG_STAR2,
     EXCLUDED_ENUM_ATTRIBUTES,
     SYMBOL_FUNCBASE_TYPES,
     Context,
@@ -359,10 +360,13 @@ def analyze_instance_member_access(
             signature = method.type
         signature = freshen_all_functions_type_vars(signature)
         if not method.is_static:
-            signature = check_self_arg(
-                signature, mx.self_type, method.is_class, mx.context, name, mx.msg
-            )
-            signature = bind_self(signature, mx.self_type, is_classmethod=method.is_class)
+            if isinstance(method, (FuncDef, OverloadedFuncDef)) and method.is_trivial_self:
+                signature = bind_self_fast(signature, mx.self_type)
+            else:
+                signature = check_self_arg(
+                    signature, mx.self_type, method.is_class, mx.context, name, mx.msg
+                )
+                signature = bind_self(signature, mx.self_type, is_classmethod=method.is_class)
         # TODO: should we skip these steps for static methods as well?
         # Since generic static methods should not be allowed.
         typ = map_instance_to_supertype(typ, method.info)
@@ -521,9 +525,11 @@ def analyze_member_var_access(
     mx.chk.warn_deprecated(v, mx.context)
 
     vv = v
+    is_trivial_self = False
     if isinstance(vv, Decorator):
         # The associated Var node of a decorator contains the type.
         v = vv.var
+        is_trivial_self = vv.func.is_trivial_self and not vv.decorators
         if mx.is_super and not mx.suppress_errors:
             validate_super_call(vv.func, mx)
 
@@ -555,7 +561,7 @@ def analyze_member_var_access(
         if mx.is_lvalue and not mx.chk.get_final_context():
             check_final_member(name, info, mx.msg, mx.context)
 
-        return analyze_var(name, v, itype, mx, implicit=implicit)
+        return analyze_var(name, v, itype, mx, implicit=implicit, is_trivial_self=is_trivial_self)
     elif isinstance(v, FuncDef):
         assert False, "Did not expect a function"
     elif isinstance(v, MypyFile):
@@ -850,14 +856,21 @@ def is_instance_var(var: Var) -> bool:
 
 
 def analyze_var(
-    name: str, var: Var, itype: Instance, mx: MemberContext, *, implicit: bool = False
+    name: str,
+    var: Var,
+    itype: Instance,
+    mx: MemberContext,
+    *,
+    implicit: bool = False,
+    is_trivial_self: bool = False,
 ) -> Type:
     """Analyze access to an attribute via a Var node.
 
     This is conceptually part of analyze_member_access and the arguments are similar.
     itype is the instance type in which attribute should be looked up
     original_type is the type of E in the expression E.var
     if implicit is True, the original Var was created as an assignment to self
+    if is_trivial_self is True, we can use fast path for bind_self().
     """
     # Found a member variable.
     original_itype = itype
@@ -904,7 +917,7 @@ def analyze_var(
             for ct in call_type.items if isinstance(call_type, UnionType) else [call_type]:
                 p_ct = get_proper_type(ct)
                 if isinstance(p_ct, FunctionLike) and not p_ct.is_type_obj():
-                    item = expand_and_bind_callable(p_ct, var, itype, name, mx)
+                    item = expand_and_bind_callable(p_ct, var, itype, name, mx, is_trivial_self)
                 else:
                     item = expand_without_binding(ct, var, itype, original_itype, mx)
                 bound_items.append(item)
@@ -938,13 +951,21 @@ def expand_without_binding(
 
 
 def expand_and_bind_callable(
-    functype: FunctionLike, var: Var, itype: Instance, name: str, mx: MemberContext
+    functype: FunctionLike,
+    var: Var,
+    itype: Instance,
+    name: str,
+    mx: MemberContext,
+    is_trivial_self: bool,
 ) -> Type:
     functype = freshen_all_functions_type_vars(functype)
     typ = get_proper_type(expand_self_type(var, functype, mx.original_type))
     assert isinstance(typ, FunctionLike)
-    typ = check_self_arg(typ, mx.self_type, var.is_classmethod, mx.context, name, mx.msg)
-    typ = bind_self(typ, mx.self_type, var.is_classmethod)
+    if is_trivial_self:
+        typ = bind_self_fast(typ, mx.self_type)
+    else:
+        typ = check_self_arg(typ, mx.self_type, var.is_classmethod, mx.context, name, mx.msg)
+        typ = bind_self(typ, mx.self_type, var.is_classmethod)
     expanded = expand_type_by_instance(typ, itype)
     freeze_all_type_vars(expanded)
     if not var.is_property:
@@ -1203,10 +1224,22 @@ def analyze_class_attribute_access(
             isinstance(node.node, SYMBOL_FUNCBASE_TYPES) and node.node.is_static
         )
         t = get_proper_type(t)
-        if isinstance(t, FunctionLike) and is_classmethod:
+        is_trivial_self = False
+        if isinstance(node.node, Decorator):
+            # Use fast path if there are trivial decorators like @classmethod or @property
+            is_trivial_self = node.node.func.is_trivial_self and not node.node.decorators
+        elif isinstance(node.node, (FuncDef, OverloadedFuncDef)):
+            is_trivial_self = node.node.is_trivial_self
+        if isinstance(t, FunctionLike) and is_classmethod and not is_trivial_self:
             t = check_self_arg(t, mx.self_type, False, mx.context, name, mx.msg)
         result = add_class_tvars(
-            t, isuper, is_classmethod, is_staticmethod, mx.self_type, original_vars=original_vars
+            t,
+            isuper,
+            is_classmethod,
+            is_staticmethod,
+            mx.self_type,
+            original_vars=original_vars,
+            is_trivial_self=is_trivial_self,
         )
         # __set__ is not called on class objects.
         if not mx.is_lvalue:
@@ -1255,7 +1288,7 @@ def analyze_class_attribute_access(
         # Annotated and/or explicit class methods go through other code paths above, for
         # unannotated implicit class methods we do this here.
         if node.node.is_class:
-            typ = bind_self(typ, is_classmethod=True)
+            typ = bind_self_fast(typ)
         return apply_class_attr_hook(mx, hook, typ)
 
 
@@ -1342,6 +1375,7 @@ def add_class_tvars(
     is_staticmethod: bool,
     original_type: Type,
     original_vars: Sequence[TypeVarLikeType] | None = None,
+    is_trivial_self: bool = False,
 ) -> Type:
     """Instantiate type variables during analyze_class_attribute_access,
     e.g T and Q in the following:
@@ -1362,6 +1396,7 @@ class B(A[str]): pass
         original_type: The value of the type B in the expression B.foo() or the corresponding
             component in case of a union (this is used to bind the self-types)
         original_vars: Type variables of the class callable on which the method was accessed
+        is_trivial_self: if True, we can use fast path for bind_self().
     Returns:
         Expanded method type with added type variables (when needed).
     """
@@ -1383,7 +1418,10 @@ class B(A[str]): pass
         tvars = original_vars if original_vars is not None else []
         t = freshen_all_functions_type_vars(t)
         if is_classmethod:
-            t = bind_self(t, original_type, is_classmethod=True)
+            if is_trivial_self:
+                t = bind_self_fast(t, original_type)
+            else:
+                t = bind_self(t, original_type, is_classmethod=True)
         if is_classmethod or is_staticmethod:
             assert isuper is not None
             t = expand_type_by_instance(t, isuper)
@@ -1422,5 +1460,45 @@ def analyze_decorator_or_funcbase_access(
     if isinstance(defn, Decorator):
         return analyze_var(name, defn.var, itype, mx)
     typ = function_type(defn, mx.chk.named_type("builtins.function"))
+    is_trivial_self = False
+    if isinstance(defn, Decorator):
+        # Use fast path if there are trivial decorators like @classmethod or @property
+        is_trivial_self = defn.func.is_trivial_self and not defn.decorators
+    elif isinstance(defn, (FuncDef, OverloadedFuncDef)):
+        is_trivial_self = defn.is_trivial_self
+    if is_trivial_self:
+        return bind_self_fast(typ, mx.self_type)
     typ = check_self_arg(typ, mx.self_type, defn.is_class, mx.context, name, mx.msg)
     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 cast(F, method)
+    if method.arg_kinds[0] in (ARG_STAR, ARG_STAR2):
+        # See typeops.py for details.
+        return cast(F, method)
+    original_type = get_proper_type(original_type)
+    if isinstance(original_type, CallableType) and original_type.is_type_obj():
+        original_type = TypeType.make_normalized(original_type.ret_type)
+    res = method.copy_modified(
+        arg_types=method.arg_types[1:],
+        arg_kinds=method.arg_kinds[1:],
+        arg_names=method.arg_names[1:],
+        bound_args=[original_type],
+    )
+    return cast(F, res)

mypy/nodes.py

@@ -550,7 +550,7 @@ class OverloadedFuncDef(FuncBase, SymbolNode, Statement):
     Overloaded variants must be consecutive in the source file.
     """
 
-    __slots__ = ("items", "unanalyzed_items", "impl", "deprecated")
+    __slots__ = ("items", "unanalyzed_items", "impl", "deprecated", "_is_trivial_self")
 
     items: list[OverloadPart]
     unanalyzed_items: list[OverloadPart]
@@ -563,6 +563,7 @@ def __init__(self, items: list[OverloadPart]) -> None:
         self.unanalyzed_items = items.copy()
         self.impl = None
         self.deprecated = None
+        self._is_trivial_self: bool | None = None
         if items:
             # TODO: figure out how to reliably set end position (we don't know the impl here).
             self.set_line(items[0].line, items[0].column)
@@ -576,6 +577,27 @@ def name(self) -> str:
             assert self.impl is not None
             return self.impl.name
 
+    @property
+    def is_trivial_self(self) -> bool:
+        """Check we can use bind_self() fast path for this overload.
+
+        This will return False if at least one overload:
+          * Has an explicit self annotation, or Self in signature.
+          * Has a non-trivial decorator.
+        """
+        if self._is_trivial_self is not None:
+            return self._is_trivial_self
+        for item in self.items:
+            if isinstance(item, FuncDef):
+                if not item.is_trivial_self:
+                    self._is_trivial_self = False
+                    return False
+            elif item.decorators or not item.func.is_trivial_self:
+                self._is_trivial_self = False
+                return False
+        self._is_trivial_self = True
+        return True
+
     def accept(self, visitor: StatementVisitor[T]) -> T:
         return visitor.visit_overloaded_func_def(self)
 
@@ -747,6 +769,7 @@ def is_dynamic(self) -> bool:
     "is_decorated",
     "is_conditional",
     "is_trivial_body",
+    "is_trivial_self",
     "is_mypy_only",
 ]
 
@@ -771,6 +794,7 @@ class FuncDef(FuncItem, SymbolNode, Statement):
         "abstract_status",
         "original_def",
         "is_trivial_body",
+        "is_trivial_self",
         "is_mypy_only",
         # Present only when a function is decorated with @typing.dataclass_transform or similar
         "dataclass_transform_spec",
@@ -804,6 +828,10 @@ def __init__(
         self.dataclass_transform_spec: DataclassTransformSpec | None = None
         self.docstring: str | None = None
         self.deprecated: str | None = None
+        # This is used to simplify bind_self() logic in trivial cases (which are
+        # the majority). In cases where self is not annotated and there are no Self
+        # in the signature we can simply drop the first argument.
+        self.is_trivial_self = False
 
     @property
     def name(self) -> str:

mypy/semanal.py

@@ -1085,6 +1085,7 @@ def prepare_method_signature(self, func: FuncDef, info: TypeInfo, has_self_type:
                         assert self.type is not None and self.type.self_type is not None
                         leading_type: Type = self.type.self_type
                     else:
+                        func.is_trivial_self = True
                         leading_type = fill_typevars(info)
                     if func.is_class or func.name == "__new__":
                         leading_type = self.class_type(leading_type)