Support python 3.8 to 3.10

cloudpickle_generators/__init__.py

@@ -221,29 +221,86 @@ def _save_generator_impl(self, frame, gen, filler):
     write(pickle.REDUCE)
 
 
+# The _reduce.*() variants work with the C implementation of pickle, which
+# became the default from Python 3.8.
+def _reduce_generator(gen):
+    if gen.gi_running:
+        raise ValueError("cannot save running generator")
+
+    frame = gen.gi_frame
+    return _reduce_generator_impl(frame, gen, _fill_generator)
+
+
+def _reduce_coroutine(coro):
+    frame = coro.cr_frame
+    return _reduce_generator_impl(frame, coro, _fill_coroutine)
+
+
+def _reduce_async_generator(asyncgen):
+    frame = asyncgen.ag_frame
+    return _reduce_generator_impl(frame, asyncgen, _fill_async_generator)
+
+
+def _reduce_generator_impl(frame, gen, filler):
+    if frame is None:
+        # frame is None when the generator is fully consumed; take a fast path
+        return _restore_spent_generator, (
+            gen.__name__,
+            getattr(gen, "__qualname__", None),
+        )
+
+    f_code = frame.f_code
+
+    # Create a copy of generator function without the closure to serve as a box
+    # to serialize the code, globals, name, and closure. Cloudpickle already
+    # handles things like closures and complicated globals so just rely on
+    # cloudpickle to serialize this function.
+    gen_func = FunctionType(
+        f_code,
+        frame.f_globals,
+        gen.__name__,
+        (),
+        (_empty_cell(),) * len(f_code.co_freevars),
+    )
+    gen_func.__qualname__ = gen.__qualname__
+
+    return (
+        _create_skeleton_generator,
+        (gen_func,),
+        (frame.f_lasti, frame.f_locals, private_frame_data(frame)),
+        None,
+        None,
+        lambda obj, state: filler(obj, *state),
+    )
+
+
 def register():
     """Register the cloudpickle extension.
     """
-    CloudPickler.dispatch[GeneratorType] = _save_generator
-    if sys.version_info >= (3, 5, 0):
-        CloudPickler.dispatch[CoroutineType] = _save_coroutine
-    if sys.version_info >= (3, 6, 0):
-        CloudPickler.dispatch[AsyncGeneratorType] = _save_async_generator
+    if pickle.HIGHEST_PROTOCOL >= 5:
+        CloudPickler.dispatch[GeneratorType] = _reduce_generator
+        CloudPickler.dispatch[CoroutineType] = _reduce_coroutine
+        CloudPickler.dispatch[AsyncGeneratorType] = _reduce_async_generator
+    else:
+        CloudPickler.dispatch[GeneratorType] = _save_generator
+        if sys.version_info >= (3, 5, 0):
+            CloudPickler.dispatch[CoroutineType] = _save_coroutine
+        if sys.version_info >= (3, 6, 0):
+            CloudPickler.dispatch[AsyncGeneratorType] = _save_async_generator
 
 
 def unregister():
     """Unregister the cloudpickle extension.
     """
-    if CloudPickler.dispatch.get(GeneratorType) is _save_generator:
+    if CloudPickler.dispatch.get(GeneratorType) in [_save_generator, _reduce_generator]:
         # make sure we are only removing the dispatch we added, not someone
         # else's
         del CloudPickler.dispatch[GeneratorType]
 
     if sys.version_info >= (3, 5, 0):
-        if CloudPickler.dispatch.get(CoroutineType) is _save_coroutine:
+        if CloudPickler.dispatch.get(CoroutineType) is [_save_coroutine, _reduce_coroutine]:
             del CloudPickler.dispatch[CoroutineType]
 
     if sys.version_info >= (3, 6, 0):
-        if (CloudPickler.dispatch.get(AsyncGeneratorType) is
-                _save_async_generator):
+        if CloudPickler.dispatch.get(AsyncGeneratorType) in [_save_async_generator, _reduce_async_generator]:
             del CloudPickler.dispatch[AsyncGeneratorType]

cloudpickle_generators/_core.c

@@ -199,7 +199,12 @@ private_frame_data(PyObject* UNUSED(self), PyObject* frame_ob) {
     }
 
     frame = (PyFrameObject*) frame_ob;
+#if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 10
     size = frame->f_stacktop - frame->f_valuestack;
+#else
+    size = frame->f_stackdepth;
+#endif
+
 
     if (!(stack = PyTuple_New(size))) {
         return NULL;
@@ -299,7 +304,12 @@ restore_frame(PyObject* UNUSED(self), PyObject* args, PyObject* kwargs) {
     }
 
     /* set the lasti to move the generator's instruction pointer */
+#if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 10
     frame->f_lasti = lasti;
+#else
+    // In python3.10, the f_lasti returned in python is different from the f_lasti in the C struct.
+    frame->f_lasti = lasti / sizeof(_Py_CODEUNIT);
+#endif
 
     /* restore the local variable state */
     for (ix = 0; ix < PyList_Size(locals); ++ix) {
@@ -317,7 +327,11 @@ restore_frame(PyObject* UNUSED(self), PyObject* args, PyObject* kwargs) {
             ob = NULL;
         }
         Py_XINCREF(ob);
+#if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 10
         *frame->f_stacktop++ = ob;
+#else
+        frame->f_valuestack[frame->f_stackdepth++] = ob;
+#endif
     }
 
     /* restore the block stack (exceptions and loops) state */

cloudpickle_generators/tests/py36/test_async_generators.py

@@ -2,6 +2,8 @@
 from types import FunctionType, coroutine
 
 import cloudpickle
+import pytest
+import sys
 
 
 @coroutine
@@ -57,6 +59,7 @@ async def genfunc():
         yield 2
 
 
+@pytest.mark.xfail(sys.version_info >= (3, 8, 0), reason="asyncgen_asgen doesn't work on this function before roundtrip")
 def test_async_generator_1():
     async def ticker(delay, to):
         # PEP525