[mypyc] Move setting of __cpyfunction__ outside of constructor

Author: p-sawicki

mypyc/codegen/emitclass.py

@@ -413,7 +413,7 @@ def emit_line() -> None:
     emitter.emit_line()
     if generate_full:
         generate_setup_for_class(
-            cl, defaults_fn, vtable_name, shadow_vtable_name, coroutine_setup_name, emitter
+            cl, defaults_fn, vtable_name, shadow_vtable_name, emitter
         )
         emitter.emit_line()
         generate_constructor_for_class(cl, cl.ctor, init_fn, setup_name, vtable_name, emitter)
@@ -606,7 +606,6 @@ 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."""
@@ -662,13 +661,6 @@ 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("}")
 

mypyc/irbuild/builder.py

@@ -69,6 +69,7 @@
     Assign,
     BasicBlock,
     Branch,
+    Call,
     ComparisonOp,
     GetAttr,
     InitStatic,
@@ -91,6 +92,7 @@
     RType,
     RUnion,
     bitmap_rprimitive,
+    bool_rprimitive,
     bytes_rprimitive,
     c_pyssize_t_rprimitive,
     dict_rprimitive,
@@ -1460,6 +1462,20 @@ def add_function(self, func_ir: FuncIR, line: int) -> None:
     def get_current_class_ir(self) -> ClassIR | None:
         type_info = self.fn_info.fitem.info
         return self.mapper.type_to_ir.get(type_info)
+    
+    def add_coroutine_setup_call(self, class_name: str, obj: Value) -> Value:
+        return self.add(
+            Call(
+                FuncDecl(
+                    class_name + "_coroutine_setup",
+                    None,
+                    self.module_name,
+                    FuncSignature([RuntimeArg("type", object_rprimitive)], bool_rprimitive),
+                ),
+                [obj],
+                -1,
+            )
+        )
 
 
 def gen_arg_defaults(builder: IRBuilder) -> None:

mypyc/irbuild/callable_class.py

@@ -232,4 +232,9 @@ def instantiate_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> Value:
         curr_env_reg = builder.fn_info.curr_env_reg
     if curr_env_reg:
         builder.add(SetAttr(func_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line))
+    # 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 not fn_info.in_non_ext and fn_info.is_coroutine:
+        builder.add_coroutine_setup_call(fn_info.callable_class.ir.name, func_reg)
     return func_reg

mypyc/irbuild/classdef.py

@@ -473,19 +473,7 @@ def allocate_class(builder: IRBuilder, cdef: ClassDef) -> Value:
                 -1,
             )
         )
-
-        builder.add(
-            Call(
-                FuncDecl(
-                    cdef.name + "_coroutine_setup",
-                    None,
-                    builder.module_name,
-                    FuncSignature([RuntimeArg("type", object_rprimitive)], bool_rprimitive),
-                ),
-                [tp],
-                -1,
-            )
-        )
+        builder.add_coroutine_setup_call(cdef.name, tp)
 
     # Populate a '__mypyc_attrs__' field containing the list of attrs
     builder.primitive_op(