Issue #112: Prepare Stackless 3.5, enable pickling of coroutines

Coroutines are a new feature of Python 3.5. Because their implementation is based on generators, we can pickle and unpickle them easily.

https://bitbucket.org/stackless-dev/stackless/issues/112

Author: Anselm Kruis

Stackless/changelog.txt

@@ -87,6 +87,7 @@ What's New in Stackless 3.X.X?
   Adapt Stackless to Python 3.5.
   - PyGenObject got two additional fields
   - Skip the CPython test case test.test_pickle.*.test_local_lookup_error
+  - Enable the pickling of coroutine objects
 
 - https://bitbucket.org/stackless-dev/stackless/issues/111
   Restore the Python ABI function PyGen_New(). Previously Stackless named this

Stackless/pickling/prickelpit.c

@@ -1688,16 +1688,20 @@ static int init_methodwrappertype(void)
 
 /******************************************************
 
-  pickling of generators
+  pickling of generators and coroutines
 
  ******************************************************/
 
 static PyTypeObject wrap_PyGen_Type;
-/* Used to initialize a generator created by gen_new. */
-static PyFrameObject *gen_exhausted_frame;
+static PyTypeObject wrap_PyCoro_Type;
+
+/* Used to initialize a generator created by gen_new. 
+   Also assert, that the size of generator and coroutines is equal. */
+static PyFrameObject *gen_exhausted_frame = \
+Py_BUILD_ASSERT_EXPR(sizeof(PyGenObject) == sizeof(PyCoroObject)); /* value is 0 */
 
 static PyObject *
-gen_reduce(PyGenObject *gen)
+_gen_reduce(PyTypeObject *type, PyGenObject *gen)
 {
     PyObject *tup;
     PyObject *frame_reducer = (PyObject *)gen->gi_frame;
@@ -1712,7 +1716,7 @@ gen_reduce(PyGenObject *gen)
     if (frame_reducer == NULL)
         return NULL;
     tup = Py_BuildValue("(O()(OiOO))",
-                        &wrap_PyGen_Type,
+                        type,
                         frame_reducer,
                         gen->gi_running,
                         gen->gi_name,
@@ -1722,6 +1726,12 @@ gen_reduce(PyGenObject *gen)
     return tup;
 }
 
+static PyObject *
+gen_reduce(PyGenObject *gen)
+{
+    return _gen_reduce(&wrap_PyGen_Type, gen);
+}
+
 static PyObject *
 gen_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 {
@@ -1731,7 +1741,13 @@ gen_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     /* A reference to frame is stolen by PyGen_New. */
     assert(gen_exhausted_frame != NULL);
     assert(PyFrame_Check(gen_exhausted_frame));
-    gen = (PyGenObject *) PyGen_NewWithQualName(slp_ensure_new_frame(gen_exhausted_frame), NULL, NULL);
+    if (type == &wrap_PyGen_Type) {
+        gen = (PyGenObject *)PyGen_NewWithQualName(slp_ensure_new_frame(gen_exhausted_frame), NULL, NULL);
+    }
+    else {
+        assert(type == &wrap_PyCoro_Type);
+        gen = (PyGenObject *)PyCoro_New(slp_ensure_new_frame(gen_exhausted_frame), NULL, NULL);
+    }
     if (gen == NULL)
         return NULL;
     Py_TYPE(gen) = type;
@@ -1897,6 +1913,22 @@ static int init_generatortype(void)
 #undef initchain
 #define initchain init_generatortype
 
+static PyObject *
+coro_reduce(PyGenObject *gen)
+{
+    return _gen_reduce(&wrap_PyCoro_Type, gen);
+}
+
+MAKE_WRAPPERTYPE(PyCoro_Type, coro, "coroutine", coro_reduce,
+    gen_new, gen_setstate)
+
+static int init_coroutinetype(void)
+{
+    return init_type(&wrap_PyCoro_Type, initchain);
+}
+#undef initchain
+#define initchain init_coroutinetype
+
 
 /******************************************************
 

Stackless/unittests/test_pickle.py

@@ -632,6 +632,60 @@ def d():
         self.assertListEqual([i[1:] for i in all_outerframes_orig[:l - 1]], [i[1:] for i in all_outerframes[:l - 1]])
 
 
+class TestCoroutinePickling(StacklessPickleTestCase):
+    @types.coroutine
+    def yield1(self, value):
+        yield value
+
+    async def c(self):
+        await self.yield1(1)
+
+    def test_without_pickling(self):
+        c = self.c()
+        self.assertIsInstance(c, types.CoroutineType)
+        self.assertIsNone(c.cr_await)
+        self.assertIsNotNone(c.cr_frame)
+        self.assertEqual(c.send(None), 1)
+        self.assertIsNotNone(c.cr_await)
+        self.assertIsNotNone(c.cr_frame)
+        self.assertRaises(StopIteration, c.send, None)
+        self.assertIsNone(c.cr_await)
+        self.assertIsNone(c.cr_frame)
+        self.assertRaisesRegex(RuntimeError, "cannot reuse already awaited coroutine", c.send, None)
+
+    def test_pickling1(self):
+        c = self.c()
+        p = self.dumps(c)
+        c.send(None)
+        c = self.loads(p)
+        self.assertIsInstance(c, types.CoroutineType)
+        self.assertIsNone(c.cr_await)
+        self.assertIsNotNone(c.cr_frame)
+        self.assertEqual(c.send(None), 1)
+        self.assertRaises(StopIteration, c.send, None)
+
+    def test_pickling2(self):
+        c = self.c()
+        self.assertEqual(c.send(None), 1)
+        p = self.dumps(c)
+        c = self.loads(p)
+        self.assertIsInstance(c, types.CoroutineType)
+        self.assertIsNotNone(c.cr_await)
+        self.assertIsNotNone(c.cr_frame)
+        self.assertRaises(StopIteration, c.send, None)
+
+    def test_pickling3(self):
+        c = self.c()
+        self.assertEqual(c.send(None), 1)
+        self.assertRaises(StopIteration, c.send, None)
+        p = self.dumps(c)
+        c = self.loads(p)
+        self.assertIsInstance(c, types.CoroutineType)
+        self.assertIsNone(c.cr_await)
+        self.assertIsNone(c.cr_frame)
+        self.assertRaisesRegex(RuntimeError, "cannot reuse already awaited coroutine", c.send, None)
+
+
 class TestCopy(StacklessTestCase):
     ITERATOR_TYPE = type(iter("abc"))
 
@@ -653,6 +707,7 @@ def _test(self, obj, *attributes, **kw):
             # it is a shallow copy, therefore the attributes should
             # refer to the same objects
             self.assertIs(value_c, value_obj)
+        return c
 
     def test_module_stackless(self):
         # test for issue 128
@@ -706,7 +761,7 @@ def test_generator(self):
         def g():
             yield 1
         obj = g()
-        self._test(obj, 'gi_running', 'gi_code')
+        self._test(obj, 'gi_running', 'gi_code', '__name__', '__qualname__')
 
     def test_dict_keys(self):
         d = {1: 10, "a": "ABC"}
@@ -723,6 +778,15 @@ def test_dict_items(self):
         obj = d.items()
         self._test(obj)
 
+    def test_coroutine(self):
+        async def c():
+            return 1
+        obj = c()
+        c = self._test(obj, 'cr_running', 'cr_code', '__name__', '__qualname__')
+        self.assertRaises(StopIteration, obj.send, None)
+        self.assertRaises(StopIteration, c.send, None)
+
+
 if __name__ == '__main__':
     if not sys.argv[1:]:
         sys.argv.append('-v')