Stackless issue python#196: pickling of coroutine_wrapper objects

Stackless can now pickle coroutine_wrapper objects.

Author: Anselm Kruis

Objects/genobject.c

@@ -2158,6 +2158,55 @@ async_gen_athrow_new(PyAsyncGenObject *gen, PyObject *args)
 
 
 #ifdef STACKLESS
+PyObject *
+slp_coro_wrapper_reduce(PyObject *o, PyTypeObject * wrapper_type)
+{
+    PyCoroWrapper *cw;
+    PyObject *tup;
+
+    assert(o != NULL);
+    assert(Py_TYPE(o) == &_PyCoroWrapper_Type);
+    assert(wrapper_type != NULL);
+
+    cw = (PyCoroWrapper *)o;
+    tup = Py_BuildValue("(O()(O))",
+        wrapper_type,
+        cw->cw_coroutine);
+    return tup;
+}
+
+PyObject *
+slp_coro_wrapper_new(PyCoroObject *co)
+{
+    return coro_await(co);
+}
+
+PyObject *
+slp_coro_wrapper_setstate(PyObject *self, PyObject *args)
+{
+    PyCoroWrapper *cw;
+    PyObject *co;
+
+    assert(self != NULL);
+    cw = (PyCoroWrapper *)self;
+
+    assert(args != NULL);
+    if (!PyArg_ParseTuple(args, "O!:coro_wrapper_setstate",
+        &PyCoro_Type,
+        &co)) {
+        return NULL;
+    }
+
+    assert(PyCoro_CheckExact(co));
+    Py_INCREF(co);
+    Py_SETREF(cw->cw_coroutine, (PyCoroObject*)co);
+
+    Py_TYPE(self) = Py_TYPE(self)->tp_base;
+    Py_INCREF(self);
+    return self;
+}
+
+
 int
 slp_async_gen_init_hooks(PyAsyncGenObject *o)
 {

Stackless/changelog.txt

@@ -9,6 +9,9 @@ What's New in Stackless 3.X.X?
 
 *Release date: 20XX-XX-XX*
 
+- https://github.com/stackless-dev/stackless/issues/196
+  Stackless can now pickle coroutine_wrapper objects.
+
 - https://github.com/stackless-dev/stackless/issues/190
   Silently ignore attempts to close a running generator, coroutine or
   asynchronous generator. This avoids spurious error messages, if such an

Stackless/core/stackless_impl.h

@@ -741,6 +741,9 @@ int slp_resurrect_and_kill(PyObject *self,
                            void(*killer)(PyObject *));
 
 /* stackless pickling support */
+PyObject * slp_coro_wrapper_reduce(PyObject *o, PyTypeObject * wrapper_type);
+PyObject * slp_coro_wrapper_new(PyCoroObject *gen);
+PyObject * slp_coro_wrapper_setstate(PyObject *self, PyObject *args);
 int slp_async_gen_init_hooks(PyAsyncGenObject *o);
 PyObject * slp_async_gen_asend_reduce(PyObject *o, PyTypeObject * wrapper_type);
 PyObject * slp_async_gen_asend_new(PyAsyncGenObject *gen);

Stackless/pickling/prickelpit.c

@@ -1520,12 +1520,14 @@ static int init_dictitemsviewtype(PyObject * mod)
 
 static PyTypeObject wrap_PyGen_Type;
 static PyTypeObject wrap_PyCoro_Type;
+static PyTypeObject wrap__PyCoroWrapper_Type;
 static PyTypeObject wrap_PyAsyncGen_Type;
 static PyTypeObject wrap__PyAsyncGenASend_Type;
 static PyTypeObject wrap__PyAsyncGenAThrow_Type;
 
 /* Used to initialize a generator created by gen_new. */
 static PyFrameObject *gen_exhausted_frame = NULL;
+static PyCoroObject *gen_exhausted_coro = NULL;
 static PyAsyncGenObject *gen_exhausted_asyncgen = NULL;
 
 /* A helper for pickling the _PyErr_StackItem* members of generator-like and tasklet
@@ -1850,8 +1852,11 @@ static int init_generatortype(PyObject * mod)
         /* A reference to frame is stolen by PyGen_New. */
         Py_INCREF(gen_exhausted_frame);
         gen_exhausted_asyncgen = (PyAsyncGenObject *)PyAsyncGen_New(gen_exhausted_frame, NULL, NULL);
+        /* A reference to frame is stolen by PyCoro_New. */
+        Py_INCREF(gen_exhausted_frame);
+        gen_exhausted_coro = (PyCoroObject *)PyCoro_New(gen_exhausted_frame, NULL, NULL);
     }
-    if (gen_exhausted_asyncgen == NULL) {
+    if (gen_exhausted_asyncgen == NULL || gen_exhausted_coro == NULL) {
         res = -1;
     }
 
@@ -2194,6 +2199,41 @@ init_async_generator_athrow_type(PyObject * mod)
 #undef initchain
 #define initchain init_async_generator_athrow_type
 
+static PyObject *
+coro_wrapper_reduce(PyObject *o)
+{
+    return slp_coro_wrapper_reduce(o, &wrap__PyCoroWrapper_Type);
+}
+
+static PyObject *
+coro_wrapper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+    PyObject * o;
+    if (is_wrong_type(type)) return NULL;
+    assert(type == &wrap__PyCoroWrapper_Type);
+    assert(gen_exhausted_coro != NULL);
+    o = slp_coro_wrapper_new(gen_exhausted_coro);
+    if (o != NULL)
+        Py_TYPE(o) = type;
+    return o;
+}
+
+static PyObject *
+coro_wrapper_setstate(PyObject *self, PyObject *args)
+{
+    if (is_wrong_type(Py_TYPE(self))) return NULL;
+    return slp_coro_wrapper_setstate(self, args);
+}
+
+MAKE_WRAPPERTYPE(_PyCoroWrapper_Type, coro_wrapper, "coroutine_wrapper", coro_wrapper_reduce,
+    coro_wrapper_new, coro_wrapper_setstate)
+
+static int init_coroutine_wrapper_type(PyObject * mod)
+{
+    return init_type(&wrap__PyCoroWrapper_Type, initchain, mod);
+}
+#undef initchain
+#define initchain init_coroutine_wrapper_type
 
 /******************************************************
 

Stackless/unittests/test_pickle.py

@@ -707,6 +707,21 @@ def test_pickling_origin(self):
         self.assertTupleEqual(orig_origin, origin)
         self.assertRaises(StopIteration, c.send, None)
 
+    def test_coro_wrapper_pickling(self):
+        c = self.c()
+        cw = c.__await__()
+        cw_type = type(cw)
+        self.assertEqual(cw_type.__name__, "coroutine_wrapper")
+        p = self.dumps(cw)
+        self.assertEqual(c.send(None), 1)
+        self.assertRaises(StopIteration, c.send, None)
+
+        cw = self.loads(p)
+        self.assertIsInstance(cw, cw_type)
+        self.assertEqual(cw.send(None), 1)
+        self.assertRaises(StopIteration, cw.send, None)
+        self.assertRaisesRegex(RuntimeError, "cannot reuse already awaited coroutine", cw.send, None)
+
 
 async_gen_finalize_counter = 0
 
@@ -1011,6 +1026,14 @@ async def c():
         self.assertRaises(StopIteration, obj.send, None)
         self.assertRaises(StopIteration, c.send, None)
 
+    def test_coro_wrapper(self):
+        async def c():
+            return 1
+        obj = c().__await__()
+        cw = self._test(obj)
+        self.assertRaises(StopIteration, cw.send, None)
+        self.assertRaisesRegex(RuntimeError, "cannot reuse already awaited coroutine", obj.send, None)
+
     def test_async_generator(self):
         async def ag():
             yield 100

Tools/c-globals/ignored-globals.txt

@@ -500,6 +500,7 @@ Py_UnwindToken
 # justification: constant singleton
 gen_exhausted_frame
 gen_exhausted_asyncgen
+gen_exhausted_coro
 
 # justification: constant
 _new_methoddef