feat: support iscoroutinefunction introspection

Author: BobTheBuidler

mypyc/codegen/emitmodule.py

@@ -954,6 +954,13 @@ def generate_module_def(self, emitter: Emitter, module_name: str, module: Module
                 emitter.emit_lines(f"if (unlikely(!{type_struct}))", "    goto fail;")
 
         emitter.emit_lines("if (CPyGlobalsInit() < 0)", "    goto fail;")
+        # Patch async native functions so they're recognized as coroutines
+        for fn in module.functions:
+            if fn.decl.is_async and fn.class_name is None:
+                emitter.emit_line(f'PyObject *func = PyObject_GetAttrString({module_static}, "{fn.decl.name}");')
+                emitter.emit_line("if (!func) goto fail;")
+                emitter.emit_line("if (!CPyPatchAsyncCode(func)) { Py_DECREF(func); goto fail; }")
+                emitter.emit_line("Py_DECREF(func);")
 
         self.generate_top_level_call(module, emitter)
 

mypyc/ir/func_ir.py

@@ -140,6 +140,7 @@ def __init__(
         is_prop_setter: bool = False,
         is_prop_getter: bool = False,
         implicit: bool = False,
+        is_async: bool = False,
     ) -> None:
         self.name = name
         self.class_name = class_name
@@ -159,6 +160,7 @@ def __init__(
         # If True, not present in the mypy AST and must be synthesized during irbuild
         # Currently only supported for property getters/setters
         self.implicit = implicit
+        self.is_async = is_async
 
         # This is optional because this will be set to the line number when the corresponding
         # FuncIR is created
@@ -204,6 +206,7 @@ def serialize(self) -> JsonDict:
             "is_prop_setter": self.is_prop_setter,
             "is_prop_getter": self.is_prop_getter,
             "implicit": self.implicit,
+            "is_async": self.is_async,
         }
 
     # TODO: move this to FuncIR?
@@ -226,6 +229,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> FuncDecl:
             data["is_prop_setter"],
             data["is_prop_getter"],
             data["implicit"],
+            data["is_async"],
         )
 
 

mypyc/irbuild/prepare.py

@@ -185,7 +185,14 @@ def prepare_func_def(
         else (FUNC_CLASSMETHOD if fdef.is_class else FUNC_NORMAL)
     )
     sig = mapper.fdef_to_sig(fdef, options.strict_dunders_typing)
-    decl = FuncDecl(fdef.name, class_name, module_name, sig, kind)
+    decl = FuncDecl(
+        fdef.name,
+        class_name,
+        module_name,
+        sig,
+        kind=kind,
+        is_async=fdef.is_coroutine,
+    )
     mapper.func_to_decl[fdef] = decl
     return decl
 

mypyc/lib-rt/pythonsupport.c

@@ -103,4 +103,20 @@ CPyLong_AsSsize_tAndOverflow_(PyObject *vv, int *overflow)
 }
 
 
+// CPy support for async functions: patch code object to include CO_COROUTINE
+PyObject* CPyPatchAsyncCode(PyObject* func) {
+    PyObject* code = PyObject_GetAttrString(func, "__code__");
+    if (!code) {
+        return NULL;
+    }
+    PyCodeObject* codeobj = (PyCodeObject*)code;
+    codeobj->co_flags |= CO_COROUTINE;
+    int res = PyObject_SetAttrString(func, "__code__", code);
+    Py_DECREF(code);
+    if (res < 0) {
+        return NULL;
+    }
+    Py_INCREF(func);
+    return func;
+}
 #endif

mypyc/test-data/run-async.test

@@ -561,3 +561,16 @@ def test_bool() -> None:
 
 [file asyncio/__init__.pyi]
 def run(x: object) -> object: ...
+
+[case testIsCoroutineFunction]
+import asyncio
+import inspect
+
+async def foo():
+    return 1
+
+def test_asyncio_iscoroutinefunction():
+    assert asyncio.iscoroutinefunction(foo) is True, "foo should be recognized as coroutine function"
+
+def test_inspect_iscoroutinefunction():
+    assert inspect.iscoroutinefunction(foo) is True, "foo should be recognized as coroutine function"