[mypyc] Generate function wrappers for each callable class instance (#20575)

p-sawicki · web-flow · commit 131f9d92da58 · 2026-01-14T15:16:46.000+01:00
When generating callable classes for async functions, currently mypyc
will generate a single `CPyFunction` wrapper for each callable class
type. This results in incorrect behavior when the callable class is used
as a wrapper because the `@wraps(fn)` decorator will set the `__name__`
attribute of the callable class to name of `fn`.

If the callable class is instantiated multiple times, the name of
previous instances will be overridden because all instances share the
same `CPyFunction` wrapper which contains the name attribute.

To fix this, the `CPyFunction` wrappers are generated for each instance
of the callable class instead of for each type. The wrappers are stored
as an attribute of the callable class and instantiated in the setup
function.
diff --git a/mypyc/codegen/emitclass.py b/mypyc/codegen/emitclass.py
@@ -28,6 +28,7 @@
 from mypyc.common import (
     BITMAP_BITS,
     BITMAP_TYPE,
+    CPYFUNCTION_NAME,
     NATIVE_PREFIX,
     PREFIX,
     REG_PREFIX,
@@ -411,7 +412,9 @@ def emit_line() -> None:
 
     emitter.emit_line()
     if generate_full:
-        generate_setup_for_class(cl, defaults_fn, vtable_name, shadow_vtable_name, emitter)
+        generate_setup_for_class(
+            cl, defaults_fn, vtable_name, shadow_vtable_name, coroutine_setup_name, emitter
+        )
         emitter.emit_line()
         generate_constructor_for_class(cl, cl.ctor, init_fn, setup_name, vtable_name, emitter)
         emitter.emit_line()
@@ -603,6 +606,7 @@ def generate_setup_for_class(
     defaults_fn: FuncIR | None,
     vtable_name: str,
     shadow_vtable_name: str | None,
+    coroutine_setup_name: str,
     emitter: Emitter,
 ) -> None:
     """Generate a native function that allocates an instance of a class."""
@@ -658,6 +662,13 @@ def generate_setup_for_class(
     if defaults_fn is not None:
         emit_attr_defaults_func_call(defaults_fn, "self", emitter)
 
+    # Initialize function wrapper for callable classes. As opposed to regular functions,
+    # each instance of a callable class needs its own wrapper because they might be instantiated
+    # inside other functions.
+    if cl.coroutine_name:
+        emitter.emit_line(f"if ({NATIVE_PREFIX}{coroutine_setup_name}((PyObject *)self) != 1)")
+        emitter.emit_line("  return NULL;")
+
     emitter.emit_line("return (PyObject *)self;")
     emitter.emit_line("}")
 
@@ -1281,27 +1292,40 @@ def generate_coroutine_setup(
     emitter.emit_line(f"{NATIVE_PREFIX}{coroutine_setup_name}(PyObject *type)")
     emitter.emit_line("{")
 
-    if not any(fn.decl.is_coroutine for fn in cl.methods.values()):
+    error_stmt = "    return 2;"
+
+    def emit_instance(fn: FuncIR, fn_name: str) -> str:
+        filepath = emitter.filepath or ""
+        return emitter.emit_cpyfunction_instance(fn, fn_name, filepath, error_stmt)
+
+    def success() -> None:
         emitter.emit_line("return 1;")
         emitter.emit_line("}")
-        return
+
+    if cl.coroutine_name:
+        # Callable class generated for a coroutine. It stores its function wrapper as an attribute.
+        wrapper_name = emit_instance(cl.methods["__call__"], cl.coroutine_name)
+        struct_name = cl.struct_name(emitter.names)
+        attr = emitter.attr(CPYFUNCTION_NAME)
+        emitter.emit_line(f"(({struct_name} *)type)->{attr} = {wrapper_name};")
+        return success()
+
+    if not any(fn.decl.is_coroutine for fn in cl.methods.values()):
+        return success()
 
     emitter.emit_line("PyTypeObject *tp = (PyTypeObject *)type;")
 
     for fn in cl.methods.values():
         if not fn.decl.is_coroutine:
             continue
 
-        filepath = emitter.filepath or ""
-        error_stmt = "    return 2;"
         name = short_id_from_name(fn.name, fn.decl.shortname, fn.line)
-        wrapper_name = emitter.emit_cpyfunction_instance(fn, name, filepath, error_stmt)
+        wrapper_name = emit_instance(fn, name)
         name_obj = f"{wrapper_name}_name"
         emitter.emit_line(f'PyObject *{name_obj} = PyUnicode_FromString("{fn.name}");')
         emitter.emit_line(f"if (unlikely(!{name_obj}))")
         emitter.emit_line(error_stmt)
         emitter.emit_line(f"if (PyDict_SetItem(tp->tp_dict, {name_obj}, {wrapper_name}) < 0)")
         emitter.emit_line(error_stmt)
 
-    emitter.emit_line("return 1;")
-    emitter.emit_line("}")
+    return success()
diff --git a/mypyc/codegen/emitmodule.py b/mypyc/codegen/emitmodule.py
@@ -1137,16 +1137,6 @@ def emit_module_exec_func(
                 name_prefix = cl.name_prefix(emitter.names)
                 emitter.emit_line(f"CPyDef_{name_prefix}_trait_vtable_setup();")
 
-                if cl.coroutine_name:
-                    fn = cl.methods["__call__"]
-                    filepath = self.source_paths[module.fullname]
-                    name = cl.coroutine_name
-                    wrapper_name = emitter.emit_cpyfunction_instance(
-                        fn, name, filepath, error_stmt
-                    )
-                    static_name = emitter.static_name(cl.name + "_cpyfunction", module.fullname)
-                    emitter.emit_line(f"{static_name} = {wrapper_name};")
-
         emitter.emit_lines("if (CPyGlobalsInit() < 0)", "    goto fail;")
 
         self.generate_top_level_call(module, emitter)
diff --git a/mypyc/common.py b/mypyc/common.py
@@ -24,6 +24,7 @@
 PROPSET_PREFIX: Final = "__mypyc_setter__"
 SELF_NAME: Final = "__mypyc_self__"
 GENERATOR_ATTRIBUTE_PREFIX: Final = "__mypyc_generator_attribute__"
+CPYFUNCTION_NAME = "__cpyfunction__"
 
 # Max short int we accept as a literal is based on 32-bit platforms,
 # so that we can just always emit the same code.
diff --git a/mypyc/irbuild/callable_class.py b/mypyc/irbuild/callable_class.py
@@ -6,20 +6,10 @@
 
 from __future__ import annotations
 
-from mypyc.common import ENV_ATTR_NAME, PROPSET_PREFIX, SELF_NAME
+from mypyc.common import CPYFUNCTION_NAME, ENV_ATTR_NAME, PROPSET_PREFIX, SELF_NAME
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
-from mypyc.ir.ops import (
-    NAMESPACE_STATIC,
-    BasicBlock,
-    Call,
-    Integer,
-    LoadStatic,
-    Register,
-    Return,
-    SetAttr,
-    Value,
-)
+from mypyc.ir.ops import BasicBlock, Call, GetAttr, Integer, Register, Return, SetAttr, Value
 from mypyc.ir.rtypes import RInstance, c_pointer_rprimitive, int_rprimitive, object_rprimitive
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.context import FuncInfo, ImplicitClass
@@ -108,6 +98,7 @@ def add_coroutine_properties(
     Needed to make introspection functions like inspect.iscoroutinefunction work.
     """
     callable_class_ir.coroutine_name = coroutine_name
+    callable_class_ir.attributes[CPYFUNCTION_NAME] = object_rprimitive
 
     properties = {
         "__name__": cpyfunction_get_name,
@@ -123,14 +114,7 @@ def add_coroutine_properties(
     }
 
     def get_func_wrapper() -> Value:
-        return builder.add(
-            LoadStatic(
-                object_rprimitive,
-                callable_class_ir.name + "_cpyfunction",
-                builder.module_name,
-                NAMESPACE_STATIC,
-            )
-        )
+        return builder.add(GetAttr(builder.self(), CPYFUNCTION_NAME, -1))
 
     for name, primitive in properties.items():
         with builder.enter_method(callable_class_ir, name, object_rprimitive, internal=True):
diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test
@@ -1449,10 +1449,18 @@ def wrap_async(fn: F) -> F:
 def wrapped(val: int) -> int:
     return val
 
+@wrap
+def wrapped2(val: int) -> int:
+    return val * 2
+
 @wrap_async
 async def wrapped_async(val: int) -> int:
     return val
 
+@wrap_async
+async def wrapped2_async(val: int) -> int:
+    return val * 2
+
 class T:
     def returns_one(self) -> int:
         return 1
@@ -1518,6 +1526,16 @@ def test_function() -> None:
     assert is_coroutine(wrapped_async)
     assert asyncio.run(wrapped_async(22)) == 44
 
+    assert getattr(wrapped, "__name__") == "wrapped", getattr(wrapped, "__name__")
+    assert getattr(wrapped2, "__name__") == "wrapped2", getattr(wrapped2, "__name__")
+
+    assert getattr(wrapped_async, "__name__") == "wrapped_async", getattr(wrapped_async, "__name__")
+    assert getattr(wrapped2_async, "__name__") == "wrapped2_async", getattr(wrapped2_async, "__name__")
+
+    setattr(wrapped_async, "__name__", "new name")
+    assert getattr(wrapped_async, "__name__") == "new name", getattr(wrapped_async, "__name__")
+    assert getattr(wrapped2_async, "__name__") == "wrapped2_async", getattr(wrapped2_async, "__name__")
+
 def test_method() -> None:
     assert not is_coroutine(T.returns_one)
     assert is_coroutine(T.returns_one_async)
