[mypyc] Wrap async functions with function-like type

mypyc/codegen/emit.py

@@ -7,19 +7,22 @@
 import textwrap
 from typing import Callable, Final
 
+from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.literals import Literals
 from mypyc.common import (
     ATTR_PREFIX,
     BITMAP_BITS,
     FAST_ISINSTANCE_MAX_SUBCLASSES,
     HAVE_IMMORTAL,
+    MODULE_PREFIX,
     NATIVE_PREFIX,
+    PREFIX,
     REG_PREFIX,
     STATIC_PREFIX,
     TYPE_PREFIX,
 )
 from mypyc.ir.class_ir import ClassIR, all_concrete_classes
-from mypyc.ir.func_ir import FuncDecl
+from mypyc.ir.func_ir import FuncDecl, FuncIR, get_text_signature
 from mypyc.ir.ops import BasicBlock, Value
 from mypyc.ir.rtypes import (
     RInstance,
@@ -169,13 +172,15 @@ def __init__(
         context: EmitterContext,
         value_names: dict[Value, str] | None = None,
         capi_version: tuple[int, int] | None = None,
+        filepath: str | None = None,
     ) -> None:
         self.context = context
         self.capi_version = capi_version or sys.version_info[:2]
         self.names = context.names
         self.value_names = value_names or {}
         self.fragments: list[str] = []
         self._indent = 0
+        self.filepath = filepath
 
     # Low-level operations
 
@@ -1207,6 +1212,24 @@ def emit_unbox_failure_with_overlapping_error_value(
         self.emit_line(failure)
         self.emit_line("}")
 
+    def emit_cpyfunction_instance(
+        self, fn: FuncIR, name: str, filepath: str, error_stmt: str
+    ) -> str:
+        module = self.static_name(fn.decl.module_name, None, prefix=MODULE_PREFIX)
+        cname = f"{PREFIX}{fn.cname(self.names)}"
+        wrapper_name = f"{cname}_wrapper"
+        cfunc = f"(PyCFunction){cname}"
+        func_flags = "METH_FASTCALL | METH_KEYWORDS"
+        doc = f"PyDoc_STR({native_function_doc_initializer(fn)})"
+
+        code_flags = "CO_COROUTINE"
+        self.emit_line(
+            f'PyObject* {wrapper_name} = CPyFunction_New({module}, "{filepath}", "{name}", {cfunc}, {func_flags}, {doc}, {fn.line}, {code_flags});'
+        )
+        self.emit_line(f"if (unlikely(!{wrapper_name}))")
+        self.emit_line(error_stmt)
+        return wrapper_name
+
 
 def c_array_initializer(components: list[str], *, indented: bool = False) -> str:
     """Construct an initializer for a C array variable.
@@ -1238,3 +1261,11 @@ def c_array_initializer(components: list[str], *, indented: bool = False) -> str
     # Multi-line result
     res.append(indent + ", ".join(current))
     return "{\n    " + ",\n    ".join(res) + "\n" + indent + "}"
+
+
+def native_function_doc_initializer(func: FuncIR) -> str:
+    text_sig = get_text_signature(func)
+    if text_sig is None:
+        return "NULL"
+    docstring = f"{text_sig}\n--\n\n"
+    return c_string_initializer(docstring.encode("ascii", errors="backslashreplace"))

mypyc/codegen/emitclass.py

@@ -7,8 +7,13 @@
 
 from mypy.nodes import ARG_STAR, ARG_STAR2
 from mypyc.codegen.cstring import c_string_initializer
-from mypyc.codegen.emit import Emitter, HeaderDeclaration, ReturnHandler
-from mypyc.codegen.emitfunc import native_function_doc_initializer, native_function_header
+from mypyc.codegen.emit import (
+    Emitter,
+    HeaderDeclaration,
+    ReturnHandler,
+    native_function_doc_initializer,
+)
+from mypyc.codegen.emitfunc import native_function_header
 from mypyc.codegen.emitwrapper import (
     generate_bin_op_wrapper,
     generate_bool_wrapper,
@@ -21,7 +26,14 @@
     generate_richcompare_wrapper,
     generate_set_del_item_wrapper,
 )
-from mypyc.common import BITMAP_BITS, BITMAP_TYPE, NATIVE_PREFIX, PREFIX, REG_PREFIX
+from mypyc.common import (
+    BITMAP_BITS,
+    BITMAP_TYPE,
+    NATIVE_PREFIX,
+    PREFIX,
+    REG_PREFIX,
+    short_id_from_name,
+)
 from mypyc.ir.class_ir import ClassIR, VTableEntries
 from mypyc.ir.func_ir import (
     FUNC_CLASSMETHOD,
@@ -240,6 +252,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
     dealloc_name = f"{name_prefix}_dealloc"
     methods_name = f"{name_prefix}_methods"
     vtable_setup_name = f"{name_prefix}_trait_vtable_setup"
+    coroutine_setup_name = f"{name_prefix}_coroutine_setup"
 
     fields: dict[str, str] = {"tp_name": f'"{name}"'}
 
@@ -347,6 +360,8 @@ def emit_line() -> None:
             shadow_vtable_name = None
         vtable_name = generate_vtables(cl, vtable_setup_name, vtable_name, emitter, shadow=False)
         emit_line()
+        generate_coroutine_setup(cl, coroutine_setup_name, module, emitter)
+        emit_line()
     if del_method:
         generate_finalize_for_class(del_method, finalize_name, emitter)
         emit_line()
@@ -391,6 +406,10 @@ def emit_line() -> None:
         )
     )
 
+    if cl.coroutine_name:
+        cpyfunction = emitter.static_name(cl.name + "_cpyfunction", module)
+        emitter.emit_line(f"static PyObject *{cpyfunction} = NULL;")
+
     emitter.emit_line()
     if generate_full:
         generate_setup_for_class(cl, defaults_fn, vtable_name, shadow_vtable_name, emitter)
@@ -1254,3 +1273,36 @@ def native_class_doc_initializer(cl: ClassIR) -> str:
         text_sig = f"{cl.name}()"
     docstring = f"{text_sig}\n--\n\n"
     return c_string_initializer(docstring.encode("ascii", errors="backslashreplace"))
+
+
+def generate_coroutine_setup(
+    cl: ClassIR, coroutine_setup_name: str, module_name: str, emitter: Emitter
+) -> None:
+    emitter.emit_line("static bool")
+    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()):
+        emitter.emit_line("return 1;")
+        emitter.emit_line("}")
+        return
+
+    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)
+        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("}")

mypyc/codegen/emitfunc.py

@@ -5,7 +5,6 @@
 from typing import Final
 
 from mypyc.analysis.blockfreq import frequently_executed_blocks
-from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import DEBUG_ERRORS, Emitter, TracebackAndGotoHandler, c_array_initializer
 from mypyc.common import (
     GENERATOR_ATTRIBUTE_PREFIX,
@@ -18,14 +17,7 @@
     TYPE_VAR_PREFIX,
 )
 from mypyc.ir.class_ir import ClassIR
-from mypyc.ir.func_ir import (
-    FUNC_CLASSMETHOD,
-    FUNC_STATICMETHOD,
-    FuncDecl,
-    FuncIR,
-    all_values,
-    get_text_signature,
-)
+from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD, FuncDecl, FuncIR, all_values
 from mypyc.ir.ops import (
     ERR_FALSE,
     NAMESPACE_MODULE,
@@ -117,14 +109,6 @@ def native_function_header(fn: FuncDecl, emitter: Emitter) -> str:
     )
 
 
-def native_function_doc_initializer(func: FuncIR) -> str:
-    text_sig = get_text_signature(func)
-    if text_sig is None:
-        return "NULL"
-    docstring = f"{text_sig}\n--\n\n"
-    return c_string_initializer(docstring.encode("ascii", errors="backslashreplace"))
-
-
 def generate_native_function(
     fn: FuncIR, emitter: Emitter, source_path: str, module_name: str
 ) -> None:

mypyc/codegen/emitmodule.py

@@ -29,13 +29,15 @@
 from mypy.util import hash_digest, json_dumps
 from mypyc.analysis.capsule_deps import find_implicit_capsule_dependencies
 from mypyc.codegen.cstring import c_string_initializer
-from mypyc.codegen.emit import Emitter, EmitterContext, HeaderDeclaration, c_array_initializer
-from mypyc.codegen.emitclass import generate_class, generate_class_reuse, generate_class_type_decl
-from mypyc.codegen.emitfunc import (
-    generate_native_function,
+from mypyc.codegen.emit import (
+    Emitter,
+    EmitterContext,
+    HeaderDeclaration,
+    c_array_initializer,
     native_function_doc_initializer,
-    native_function_header,
 )
+from mypyc.codegen.emitclass import generate_class, generate_class_reuse, generate_class_type_decl
+from mypyc.codegen.emitfunc import generate_native_function, native_function_header
 from mypyc.codegen.emitwrapper import (
     generate_legacy_wrapper_function,
     generate_wrapper_function,
@@ -566,7 +568,7 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
 
         for module_name, module in self.modules.items():
             if multi_file:
-                emitter = Emitter(self.context)
+                emitter = Emitter(self.context, filepath=self.source_paths[module_name])
                 emitter.emit_line(f'#include "__native{self.short_group_suffix}.h"')
                 emitter.emit_line(f'#include "__native_internal{self.short_group_suffix}.h"')
 
@@ -986,6 +988,9 @@ def emit_module_methods(
         for fn in module.functions:
             if fn.class_name is not None or fn.name == TOP_LEVEL_NAME:
                 continue
+            # Coroutines are added to the module dict when the module is initialized.
+            if fn.decl.is_coroutine:
+                continue
             name = short_id_from_name(fn.name, fn.decl.shortname, fn.line)
             if is_fastcall_supported(fn, emitter.capi_version):
                 flag = "METH_FASTCALL"
@@ -1024,6 +1029,30 @@ def emit_module_def_struct(
         emitter.emit_line("};")
         emitter.emit_line()
 
+    def emit_coroutine_wrappers(self, emitter: Emitter, module: ModuleIR, globals: str) -> None:
+        """Emit insertion of coroutines into the module dict when the module is initialized.
+        Coroutines are wrapped in CPyFunction objects to enable introspection by functions like
+        inspect.iscoroutinefunction(fn).
+        """
+        for fn in module.functions:
+            if fn.class_name is not None or fn.name == TOP_LEVEL_NAME:
+                continue
+            if not fn.decl.is_coroutine:
+                continue
+
+            filepath = self.source_paths[module.fullname]
+            error_stmt = "    goto fail;"
+            name = short_id_from_name(fn.name, fn.decl.shortname, fn.line)
+            wrapper_name = emitter.emit_cpyfunction_instance(fn, name, filepath, error_stmt)
+            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({globals}, {name_obj}, (PyObject *){wrapper_name}) < 0)"
+            )
+            emitter.emit_line(error_stmt)
+
     def emit_module_exec_func(
         self, emitter: Emitter, module_name: str, module_prefix: str, module: ModuleIR
     ) -> None:
@@ -1071,20 +1100,33 @@ def emit_module_exec_func(
                 "    goto fail;",
             )
 
+        self.emit_coroutine_wrappers(emitter, module, module_globals)
+
         # HACK: Manually instantiate generated classes here
         type_structs: list[str] = []
         for cl in module.classes:
             type_struct = emitter.type_struct_name(cl)
             type_structs.append(type_struct)
             if cl.is_generated:
+                error_stmt = "    goto fail;"
                 emitter.emit_lines(
                     "{t} = (PyTypeObject *)CPyType_FromTemplate("
                     "(PyObject *){t}_template, NULL, modname);".format(t=type_struct)
                 )
-                emitter.emit_lines(f"if (unlikely(!{type_struct}))", "    goto fail;")
+                emitter.emit_lines(f"if (unlikely(!{type_struct}))", error_stmt)
                 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)

mypyc/common.py

@@ -82,6 +82,7 @@
     "misc_ops.c",
     "generic_ops.c",
     "pythonsupport.c",
+    "function_wrapper.c",
 ]
 
 # Python 3.12 introduced immortal objects, specified via a special reference count

mypyc/ir/class_ir.py

@@ -223,6 +223,9 @@ def __init__(
         # Is this a class inheriting from enum.Enum? Such classes can be special-cased.
         self.is_enum = False
 
+        # Name of the function if this a callable class representing a coroutine.
+        self.coroutine_name: str | None = None
+
     def __repr__(self) -> str:
         return (
             "ClassIR("
@@ -424,6 +427,7 @@ def serialize(self) -> JsonDict:
             "env_user_function": self.env_user_function.id if self.env_user_function else None,
             "reuse_freed_instance": self.reuse_freed_instance,
             "is_enum": self.is_enum,
+            "is_coroutine": self.coroutine_name,
         }
 
     @classmethod
@@ -481,6 +485,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
         )
         ir.reuse_freed_instance = data["reuse_freed_instance"]
         ir.is_enum = data["is_enum"]
+        ir.coroutine_name = data["is_coroutine"]
 
         return ir