Stackless issue #139: pickle PyAsyncGenASend or PyAsyncGenAThrow

Stackless now pickles awaitable objects of type PyAsyncGenASend or
PyAsyncGenAThrow (see PEP 525).

Author: Anselm Kruis

Doc/library/stackless/pickling.rst

@@ -116,14 +116,15 @@ In order to be able to pickle tasklets |SLP| needs to be able to pickle
 several other objects, which can't be pickled by |CPY|. |SLP|
 uses :func:`copyreg.pickle` to register “reduction” functions for the following
 types:
-:data:`~types.FunctionType`,
 :data:`~types.AsyncGeneratorType`,
 :data:`~types.CodeType`,
 :data:`~types.CoroutineType`,
+:data:`~types.FunctionType`,
 :data:`~types.GeneratorType`,
 :data:`~types.ModuleType`,
 :data:`~types.TracebackType`,
-:ref:`Cell Objects <cell-objects>` and
+:ref:`Cell Objects <cell-objects>`,
+C-types PyAsyncGenASend and PyAsyncGenAThrow (see :pep:`525`) as well as
 all kinds of :ref:`Dictionary view objects <dict-views>`.
 
 Frames

Objects/genobject.c

@@ -1465,15 +1465,6 @@ async_gen_init_hooks(PyAsyncGenObject *o)
 }
 
 
-#ifdef STACKLESS
-int
-slp_async_gen_init_hooks(PyAsyncGenObject *o)
-{
-    return async_gen_init_hooks(o);
-}
-#endif
-
-
 static PyObject *
 async_gen_anext(PyAsyncGenObject *o)
 {
@@ -2199,3 +2190,140 @@ async_gen_athrow_new(PyAsyncGenObject *gen, PyObject *args)
     _PyObject_GC_TRACK((PyObject*)o);
     return (PyObject*)o;
 }
+
+
+#ifdef STACKLESS
+int
+slp_async_gen_init_hooks(PyAsyncGenObject *o)
+{
+    return async_gen_init_hooks(o);
+}
+
+
+PyObject *
+slp_async_gen_asend_reduce(PyObject *o, PyTypeObject * wrapper_type)
+{
+    PyAsyncGenASend *asend;
+    PyObject *tup;
+
+    assert(o != NULL);
+    assert(PyAsyncGenASend_CheckExact(o));
+    assert(wrapper_type != NULL);
+
+    asend = (PyAsyncGenASend *)o;
+    tup = Py_BuildValue("(O()(OOiO))",
+                wrapper_type,
+                asend->ags_gen,
+                asend->ags_sendval ? asend->ags_sendval : Py_None,
+                (int)asend->ags_state,
+                asend->ags_sendval ? Py_True : Py_False);
+    return tup;
+}
+
+
+PyObject *
+slp_async_gen_asend_new(PyAsyncGenObject *gen)
+{
+    return async_gen_asend_new(gen, NULL);
+}
+
+PyObject *
+slp_async_gen_asend_setstate(PyObject *self, PyObject *args)
+{
+    PyAsyncGenASend *asend;
+    PyObject *gen, *send_val, *has_sendval;
+    int state;
+
+    assert(self != NULL);
+    asend = (PyAsyncGenASend *)self;
+
+    assert(args != NULL);
+    if (!PyArg_ParseTuple(args, "O!OiO!:async_gen_asend_setstate",
+        &PyAsyncGen_Type,
+        &gen,
+        &send_val,
+        &state,
+        &PyBool_Type,
+        &has_sendval)) {
+        return NULL;
+    }
+
+    assert(PyAsyncGen_CheckExact(gen));
+    Py_INCREF(gen);
+    Py_SETREF(asend->ags_gen, (PyAsyncGenObject*)gen);
+    if (has_sendval != Py_True)
+        send_val = NULL;
+    else
+        Py_INCREF(send_val);
+    Py_XSETREF(asend->ags_sendval, send_val);
+    asend->ags_state = state;
+
+    Py_TYPE(self) = Py_TYPE(self)->tp_base;
+    Py_INCREF(self);
+    return self;
+}
+
+
+PyObject *
+slp_async_gen_athrow_reduce(PyObject *o, PyTypeObject * wrapper_type)
+{
+    PyAsyncGenAThrow *athrow;
+    PyObject *tup;
+
+    assert(o != NULL);
+    assert((Py_TYPE(o) == &_PyAsyncGenAThrow_Type));
+    assert(wrapper_type != NULL);
+
+    athrow = (PyAsyncGenAThrow *)o;
+    tup = Py_BuildValue("(O()(OOiO))",
+                wrapper_type,
+                athrow->agt_gen,
+                athrow->agt_args ? athrow->agt_args : Py_None,
+                (int)athrow->agt_state,
+                athrow->agt_args ? Py_True : Py_False);
+    return tup;
+}
+
+
+PyObject *
+slp_async_gen_athrow_new(PyAsyncGenObject *gen)
+{
+    return async_gen_athrow_new(gen, NULL);
+}
+
+PyObject *
+slp_async_gen_athrow_setstate(PyObject *self, PyObject *args)
+{
+    PyAsyncGenAThrow *athrow;
+    PyObject *gen, *agt_args, *has_args;
+    int state;
+
+    assert(self != NULL);
+    athrow = (PyAsyncGenAThrow *)self;
+
+    assert(args != NULL);
+    if (!PyArg_ParseTuple(args, "O!OiO!:async_gen_athrow_setstate",
+        &PyAsyncGen_Type,
+        &gen,
+        &agt_args,
+        &state,
+        &PyBool_Type,
+        &has_args)) {
+        return NULL;
+    }
+
+    assert(PyAsyncGen_CheckExact(gen));
+    Py_INCREF(gen);
+    Py_SETREF(athrow->agt_gen, (PyAsyncGenObject*)gen);
+    if (has_args != Py_True)
+        agt_args = NULL;
+    else
+        Py_INCREF(agt_args);
+    Py_XSETREF(athrow->agt_args, agt_args);
+    athrow->agt_state = state;
+
+    Py_TYPE(self) = Py_TYPE(self)->tp_base;
+    Py_INCREF(self);
+    return self;
+}
+#endif

Stackless/changelog.txt

@@ -16,7 +16,8 @@ What's New in Stackless 3.X.X?
 - https://github.com/stackless-dev/stackless/issues/139
   New functions stackless.pickle_flags() and stackless.pickle_flags_default().
   They can be used to control pickling of certain objects.
-  Stackless can now pickle asynchronous generators.
+  Stackless can now pickle asynchronous generators and awaitables of type
+  PyAsyncGenASend or PyAsyncGenAThrow (see PEP 525).
 
 - https://github.com/stackless-dev/stackless/issues/175
   Cleanup of the Stackless C-API: Including stackless_api.h in an extension

Stackless/core/stackless_impl.h

@@ -803,6 +803,12 @@ int slp_safe_pickling(int(*save)(PyObject *, PyObject *, int),
 
 PyObject * PyStackless_Pickle_ModuleDict(PyObject *pickler, PyObject *self);
 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);
+PyObject * slp_async_gen_asend_setstate(PyObject *self, PyObject *args);
+PyObject * slp_async_gen_athrow_reduce(PyObject *o, PyTypeObject * wrapper_type);
+PyObject * slp_async_gen_athrow_new(PyAsyncGenObject *gen);
+PyObject * slp_async_gen_athrow_setstate(PyObject *self, PyObject *args);
 
 /* debugging/monitoring */
 

Stackless/pickling/prickelpit.c

@@ -1522,9 +1522,12 @@ static int init_dictitemsviewtype(PyObject * mod)
 static PyTypeObject wrap_PyGen_Type;
 static PyTypeObject wrap_PyCoro_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 PyAsyncGenObject *gen_exhausted_asyncgen = NULL;
 
 /* A helper for pickling the _PyErr_StackItem* members of generator-like and tasklet
  * objects. This method returns a pointer to the object, that contains the
@@ -1844,6 +1847,13 @@ static int init_generatortype(PyObject * mod)
     gen_exhausted_frame = slp_ensure_new_frame(gen->gi_frame);
     if (gen_exhausted_frame == NULL) {
         res = -1;
+    } else {
+        /* 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);
+    }
+    if (gen_exhausted_asyncgen == NULL) {
+        res = -1;
     }
 
     Py_DECREF(gen);
@@ -2111,6 +2121,81 @@ init_async_gentype(PyObject * mod)
 #undef initchain
 #define initchain init_async_gentype
 
+static PyObject *
+async_generator_asend_reduce(PyObject *o)
+{
+    return slp_async_gen_asend_reduce(o, &wrap__PyAsyncGenASend_Type);
+}
+
+static PyObject *
+async_generator_asend_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+    PyObject * o;
+    if (is_wrong_type(type)) return NULL;
+    assert(type == &wrap__PyAsyncGenASend_Type);
+    assert(gen_exhausted_asyncgen != NULL);
+    o = slp_async_gen_asend_new(gen_exhausted_asyncgen);
+    if (o != NULL)
+        Py_TYPE(o) = type;
+    return o;
+}
+
+static PyObject *
+async_generator_asend_setstate(PyObject *self, PyObject *args)
+{
+    if (is_wrong_type(Py_TYPE(self))) return NULL;
+    return slp_async_gen_asend_setstate(self, args);
+}
+
+MAKE_WRAPPERTYPE(_PyAsyncGenASend_Type, asyncgen_asend, "async_generator_asend", async_generator_asend_reduce,
+        async_generator_asend_new, async_generator_asend_setstate)
+
+static int
+init_async_generator_asend_type(PyObject * mod)
+{
+    return init_type(&wrap__PyAsyncGenASend_Type, initchain, mod);
+}
+#undef initchain
+#define initchain init_async_generator_asend_type
+
+static PyObject *
+async_generator_athrow_reduce(PyObject *o)
+{
+    return slp_async_gen_athrow_reduce(o, &wrap__PyAsyncGenAThrow_Type);
+}
+
+static PyObject *
+async_generator_athrow_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+    PyObject * o;
+    if (is_wrong_type(type)) return NULL;
+    assert(type == &wrap__PyAsyncGenAThrow_Type);
+    assert(gen_exhausted_asyncgen != NULL);
+    o = slp_async_gen_athrow_new(gen_exhausted_asyncgen);
+    if (o != NULL)
+        Py_TYPE(o) = type;
+    return o;
+}
+
+static PyObject *
+async_generator_athrow_setstate(PyObject *self, PyObject *args)
+{
+    if (is_wrong_type(Py_TYPE(self))) return NULL;
+    return slp_async_gen_athrow_setstate(self, args);
+}
+
+MAKE_WRAPPERTYPE(_PyAsyncGenAThrow_Type, asyncgen_athrow, "async_generator_athrow", async_generator_athrow_reduce,
+        async_generator_athrow_new, async_generator_athrow_setstate)
+
+static int
+init_async_generator_athrow_type(PyObject * mod)
+{
+    return init_type(&wrap__PyAsyncGenAThrow_Type, initchain, mod);
+}
+#undef initchain
+#define initchain init_async_generator_athrow_type
+
+
 /******************************************************
 
   support code for module dict pickling