patch "asyncio" and "inspect" stdlib modules to make them recognise Cython's generator type

CHANGES.rst

@@ -15,6 +15,12 @@ Features added
 
 * Tracing is supported in ``nogil`` functions/sections and module init code.
 
+* When generators are used in a Cython module and the module imports the
+  modules "inspect" and/or "asyncio", Cython enables interoperability by
+  patching these modules to recognise Cython's internal generator type.
+  This can be disabled by C compiling the module with
+  "-D CYTHON_PATCH_ASYNCIO=0" or "-D CYTHON_PATCH_INSPECT=0"
+
 * Adding/subtracting/dividing/modulus and equality comparisons with
   constant Python floats and small integers are faster.
 

Cython/Compiler/ExprNodes.py

@@ -25,7 +25,7 @@ from .Code import UtilityCode, TempitaUtilityCode
 from . import StringEncoding
 from . import Naming
 from . import Nodes
-from .Nodes import Node
+from .Nodes import Node, utility_code_for_imports
 from . import PyrexTypes
 from .PyrexTypes import py_object_type, c_long_type, typecast, error_type, \
     unspecified_type
@@ -2251,7 +2251,6 @@ class ImportNode(ExprNode):
             name_list = self.name_list.analyse_types(env)
             self.name_list = name_list.coerce_to_pyobject(env)
         self.is_temp = 1
-        env.use_utility_code(UtilityCode.load_cached("Import", "ImportExport.c"))
         return self
 
     gil_message = "Python import"
@@ -2261,13 +2260,24 @@ class ImportNode(ExprNode):
             name_list_code = self.name_list.py_result()
         else:
             name_list_code = "0"
-        code.putln(
-            "%s = __Pyx_Import(%s, %s, %d); %s" % (
-                self.result(),
-                self.module_name.py_result(),
-                name_list_code,
-                self.level,
-                code.error_goto_if_null(self.result(), self.pos)))
+
+        code.globalstate.use_utility_code(UtilityCode.load_cached("Import", "ImportExport.c"))
+        import_code = "__Pyx_Import(%s, %s, %d)" % (
+            self.module_name.py_result(),
+            name_list_code,
+            self.level)
+
+        if (self.level <= 0 and
+                self.module_name.is_string_literal and
+                self.module_name.value in utility_code_for_imports):
+            helper_func, code_name, code_file = utility_code_for_imports[self.module_name.value]
+            code.globalstate.use_utility_code(UtilityCode.load_cached(code_name, code_file))
+            import_code = '%s(%s)' % (helper_func, import_code)
+
+        code.putln("%s = %s; %s" % (
+            self.result(),
+            import_code,
+            code.error_goto_if_null(self.result(), self.pos)))
         code.put_gotref(self.py_result())
 
 

Cython/Compiler/Nodes.py

@@ -7188,6 +7188,14 @@ utility_code_for_cimports = {
 }
 
 
+utility_code_for_imports = {
+    # utility code used when special modules are imported.
+    # TODO: Consider a generic user-level mechanism for importing
+    'asyncio': ("__Pyx_patch_asyncio", "PatchAsyncIO", "Generator.c"),
+    'inspect': ("__Pyx_patch_inspect", "PatchInspect", "Generator.c"),
+}
+
+
 class CImportStatNode(StatNode):
     #  cimport statement
     #

Cython/Utility/Generator.c

@@ -739,3 +739,150 @@ static int __pyx_Generator_init(void) {
     }
     return 0;
 }
+
+
+//////////////////// PatchModuleWithGenerator.proto ////////////////////
+
+#ifdef __Pyx_Generator_USED
+static PyObject* __Pyx_Generator_patch_module(PyObject* module, const char* py_code); /*proto*/
+#endif
+
+//////////////////// PatchModuleWithGenerator ////////////////////
+//@substitute: naming
+
+#ifdef __Pyx_Generator_USED
+static PyObject* __Pyx_Generator_patch_module(PyObject* module, const char* py_code) {
+    PyObject *globals, *result_obj;
+    globals = PyDict_New();  if (unlikely(!globals)) goto ignore;
+    if (unlikely(PyDict_SetItemString(globals, "_cython_generator_type", (PyObject*)__pyx_GeneratorType) < 0)) goto ignore;
+    if (unlikely(PyDict_SetItemString(globals, "_module", module) < 0)) goto ignore;
+    if (unlikely(PyDict_SetItemString(globals, "_b", $builtins_cname) < 0)) goto ignore;
+    result_obj = PyRun_String(py_code, Py_file_input, globals, globals);
+    if (unlikely(!result_obj)) goto ignore;
+    Py_DECREF(result_obj);
+    Py_DECREF(globals);
+    return module;
+
+ignore:
+    PyErr_WriteUnraisable(module);
+    Py_XDECREF(globals);
+    if (unlikely(PyErr_WarnEx(PyExc_RuntimeWarning, "Cython module failed to patch module with custom type", 1) < 0)) {
+        Py_DECREF(module);
+        module = NULL;
+    }
+    return module;
+}
+#endif
+
+//////////////////// PatchAsyncIO.proto ////////////////////
+
+// run after importing "asyncio" to patch Cython generator support into it
+#if defined(__Pyx_Generator_USED) && (!defined(CYTHON_PATCH_ASYNCIO) || CYTHON_PATCH_ASYNCIO)
+static PyObject* __Pyx_patch_asyncio(PyObject* module); /*proto*/
+#else
+#define __Pyx_patch_asyncio(module)  (module)
+#endif
+
+//////////////////// PatchAsyncIO ////////////////////
+//@requires: PatchModuleWithGenerator
+//@requires: PatchInspect
+
+#if defined(__Pyx_Generator_USED) && (!defined(CYTHON_PATCH_ASYNCIO) || CYTHON_PATCH_ASYNCIO)
+static PyObject* __Pyx_patch_asyncio(PyObject* module) {
+    PyObject *patch_module = NULL;
+    static int asyncio_patched = 0;
+    if (unlikely((!asyncio_patched) && module)) {
+        PyObject *package;
+        package = __Pyx_Import(PYIDENT("asyncio.coroutines"), NULL, 0);
+        if (package) {
+            patch_module = __Pyx_Generator_patch_module(
+                PyObject_GetAttrString(package, "coroutines"),
+                "old_types = _b.getattr(_module, '_COROUTINE_TYPES', None)\n"
+                "if old_types is not None and _cython_generator_type not in old_types:\n"
+                "    _module._COROUTINE_TYPES = _b.type(old_types) (_b.tuple(old_types) + (_cython_generator_type,))\n"
+            );
+        #if PY_VERSION_HEX < 0x03050000
+        } else {
+            // Py3.4 used to have asyncio.tasks instead of asyncio.coroutines
+            PyErr_Clear();
+            package = __Pyx_Import(PYIDENT("asyncio.tasks"), NULL, 0);
+            if (unlikely(!package)) goto asyncio_done;
+            patch_module = __Pyx_Generator_patch_module(
+                PyObject_GetAttrString(package, "tasks"),
+                "if (_b.hasattr(_module, 'iscoroutine') and"
+                "        _b.getattr(_module.iscoroutine, '_cython_generator_type', None) is not _cython_generator_type):\n"
+                "    def cy_wrap(orig_func, cython_generator_type=_cython_generator_type, type=_b.type):\n"
+                "        def cy_iscoroutine(obj): return type(obj) is cython_generator_type or orig_func(obj)\n"
+                "        cy_iscoroutine._cython_generator_type = cython_generator_type\n"
+                "        return cy_iscoroutine\n"
+                "    _module.iscoroutine = cy_wrap(_module.iscoroutine)\n"
+            );
+        #endif
+        }
+        Py_DECREF(package);
+        if (unlikely(!patch_module)) goto ignore;
+#if PY_VERSION_HEX < 0x03050000
+asyncio_done:
+        PyErr_Clear();
+#endif
+        asyncio_patched = 1;
+        // now patch inspect.isgenerator() by looking up the imported module in the patched asyncio module
+        {
+            PyObject *inspect_module;
+            if (patch_module) {
+                inspect_module = PyObject_GetAttrString(patch_module, "inspect");
+                Py_DECREF(patch_module);
+            } else {
+                inspect_module = __Pyx_Import(PYIDENT("inspect"), NULL, 0);
+            }
+            if (unlikely(!inspect_module)) goto ignore;
+            inspect_module = __Pyx_patch_inspect(inspect_module);
+            if (unlikely(!inspect_module)) {
+                Py_DECREF(module);
+                module = NULL;
+            }
+            Py_DECREF(inspect_module);
+        }
+    }
+    return module;
+ignore:
+    PyErr_WriteUnraisable(module);
+    if (unlikely(PyErr_WarnEx(PyExc_RuntimeWarning, "Cython module failed to patch asyncio package with custom generator type", 1) < 0)) {
+        Py_DECREF(module);
+        module = NULL;
+    }
+    return module;
+}
+#endif
+
+
+//////////////////// PatchInspect.proto ////////////////////
+
+// run after importing "inspect" to patch Cython generator support into it
+#if defined(__Pyx_Generator_USED) && (!defined(CYTHON_PATCH_INSPECT) || CYTHON_PATCH_INSPECT)
+static PyObject* __Pyx_patch_inspect(PyObject* module); /*proto*/
+#else
+#define __Pyx_patch_inspect(module)  (module)
+#endif
+
+//////////////////// PatchInspect ////////////////////
+//@requires: PatchModuleWithGenerator
+
+#if defined(__Pyx_Generator_USED) && (!defined(CYTHON_PATCH_INSPECT) || CYTHON_PATCH_INSPECT)
+static PyObject* __Pyx_patch_inspect(PyObject* module) {
+    static int inspect_patched = 0;
+    if (unlikely((!inspect_patched) && module)) {
+        module = __Pyx_Generator_patch_module(
+            module,
+            "if _b.getattr(_module.isgenerator, '_cython_generator_type', None) is not _cython_generator_type:\n"
+            "    def cy_wrap(orig_func, cython_generator_type=_cython_generator_type, type=_b.type):\n"
+            "        def cy_isgenerator(obj): return type(obj) is cython_generator_type or orig_func(obj)\n"
+            "        cy_isgenerator._cython_generator_type = cython_generator_type\n"
+            "        return cy_isgenerator\n"
+            "    _module.isgenerator = cy_wrap(_module.isgenerator)\n"
+        );
+        inspect_patched = 1;
+    }
+    return module;
+}
+#endif

runtests.py

@@ -98,14 +98,15 @@ def get_distutils_distro(_cache=[]):
 
 
 EXT_DEP_MODULES = {
-    'tag:numpy' : 'numpy',
-    'tag:pstats': 'pstats',
-    'tag:posix' : 'posix',
-    'tag:array' : 'array',
+    'tag:numpy':    'numpy',
+    'tag:asyncio':  'asyncio',
+    'tag:pstats':   'pstats',
+    'tag:posix':    'posix',
+    'tag:array':    'array',
     'tag:coverage': 'Cython.Coverage',
-    'Coverage': 'Cython.Coverage',
-    'tag:ipython': 'IPython',
-    'tag:jedi': 'jedi',
+    'Coverage':     'Cython.Coverage',
+    'tag:ipython':  'IPython',
+    'tag:jedi':     'jedi',
 }
 
 def patch_inspect_isfunction():

tests/run/asyncio_generators.srctree

+# mode: run
+# tag: asyncio
+
+"""
+PYTHON setup.py build_ext -i
+PYTHON test_from_import.py
+PYTHON test_import.py
+PYTHON test_both.py
+"""
+
+######## setup.py ########
+
+from Cython.Build.Dependencies import cythonize
+from distutils.core import setup
+
+setup(
+  ext_modules = cythonize("*.pyx"),
+)
+
+
+######## test_from_import.py ########
+
+import from_asyncio_import
+import asyncio
+
+def runloop(task):
+    loop = asyncio.get_event_loop()
+    result = loop.run_until_complete(task())
+    assert 3 == result, result
+
+runloop(from_asyncio_import.wait3)
+
+
+######## test_import.py ########
+
+import import_asyncio
+import asyncio
+
+def runloop(task):
+    loop = asyncio.get_event_loop()
+    result = loop.run_until_complete(task())
+    assert 3 == result, result
+
+runloop(import_asyncio.wait3)
+
+
+######## test_both.py ########
+
+import asyncio
+
+def runloop(task):
+    loop = asyncio.get_event_loop()
+    result = loop.run_until_complete(task())
+    assert 3 == result, result
+
+import import_asyncio
+runloop(import_asyncio.wait3)       # 1a)
+import from_asyncio_import
+runloop(from_asyncio_import.wait3)  # 1b)
+
+runloop(from_asyncio_import.wait3)  # 2a)
+runloop(import_asyncio.wait3)       # 2b)
+runloop(from_asyncio_import.wait3)  # 3a)
+runloop(import_asyncio.wait3)       # 3b)
+
+
+######## import_asyncio.pyx ########
+# cython: binding=True
+
+import asyncio
+
+@asyncio.coroutine
+def wait3():
+    counter = 0
+    for i in range(3):
+        print(counter)
+        yield from asyncio.sleep(0.01)
+        counter += 1
+    return counter
+
+
+######## from_asyncio_import.pyx ########
+# cython: binding=True
+
+from asyncio import coroutine, sleep
+
+@coroutine
+def wait3():
+    counter = 0
+    for i in range(3):
+        print(counter)
+        yield from sleep(0.01)
+        counter += 1
+    return counter