Issue #130: fix pickling of exhausted generator objects

Pickling of an exhausted generator no longer causes a segmentation fault.

https://bitbucket.org/stackless-dev/stackless/issues/130
(grafted from b86bda31f4e3372bbde61179cc890af24d83ae71)

Author: hg

Stackless/changelog.txt

@@ -9,6 +9,9 @@ What's New in Stackless 3.X.X?
 
 *Release date: 20XX-XX-XX*
 
+- https://bitbucket.org/stackless-dev/stackless/issues/130
+  Pickling exhausted generators no longer crashes.
+
 - https://bitbucket.org/stackless-dev/stackless/issues/129
   C-API: Calling PyTasklet_New( NULL, ...) no longer crashes.
 

Stackless/pickling/prickelpit.c

@@ -1707,9 +1707,15 @@ static PyObject *
 gen_reduce(PyGenObject *gen)
 {
     PyObject *tup;
+    PyObject *gi_frame = (PyObject *)gen->gi_frame;
+    if (gi_frame == NULL) {
+        /* Pickle NULL as None. See gen_setstate() for the corresponding
+         * unpickling code. */
+        gi_frame = Py_None;
+    }
     tup = Py_BuildValue("(O()(Oi))",
                         &wrap_PyGen_Type,
-                        gen->gi_frame,
+                        gi_frame,
                         gen->gi_running
                         );
     return tup;
@@ -1736,12 +1742,39 @@ gen_setstate(PyObject *self, PyObject *args)
 {
     PyGenObject *gen = (PyGenObject *) self;
     PyFrameObject *f;
+    PyObject *obj;
     int gi_running;
 
     if (is_wrong_type(Py_TYPE(self))) return NULL;
-    if (!PyArg_ParseTuple(args, "O!i:generator",
-                          &PyFrame_Type, &f, &gi_running))
+    if (!PyArg_ParseTuple(args, "Oi:generator",
+                          &obj, &gi_running))
+        return NULL;
+
+    if (obj == Py_None) {
+        /* No frame, generator is exhausted */
+        Py_CLEAR(gen->gi_frame);
+
+        /* Even if gi_frame is NULL, gi_code is still valid. Therefore
+         * I set it to the code of the exhausted frame singleton. 
+         */
+        assert(gen_exhausted_frame != NULL);
+        assert(PyFrame_Check(gen_exhausted_frame));
+        obj = (PyObject *)gen_exhausted_frame->f_code;
+        Py_INCREF(obj);
+        Py_SETREF(gen->gi_code, obj);
+
+        gen->gi_running = gi_running;
+        Py_TYPE(gen) = Py_TYPE(gen)->tp_base;
+        Py_INCREF(gen);
+        return (PyObject *)gen;
+    }
+    if (!PyFrame_Check(obj)) {
+        PyErr_SetString(PyExc_TypeError,
+                        "invalid frame object for gen_setstate");
         return NULL;
+    }
+    f = (PyFrameObject *)obj;
+    obj = NULL;
 
     if (!gi_running) {
         if ((f = slp_ensure_new_frame(f)) != NULL) {

Stackless/unittests/test_generator.py

@@ -3,6 +3,7 @@
 import gc
 import stackless
 import types
+import pickle
 
 from support import test_main  # @UnusedImport
 from support import StacklessTestCase
@@ -51,5 +52,28 @@ def test_wrap_generator_frame_code(self):
         self.assertIsNot(g0.gi_frame, g1.gi_frame)
         self.assertEqual(g0.__name__, "exhausted_generator")
 
+
+class TestGeneratorPickling(StacklessTestCase):
+    def test_name(self):
+        def g():
+            yield 1
+        gen_orig = g()
+        p = pickle.dumps(gen_orig)
+        gen_new = pickle.loads(p)
+
+        self.assertEqual(gen_new.__name__, gen_orig.__name__)
+
+    def test_exhausted(self):
+        def g():
+            yield 1
+        gen_orig = g()
+        self.assertEqual(gen_orig.__next__(), 1)
+        self.assertRaises(StopIteration, gen_orig.__next__)
+        p = pickle.dumps(gen_orig)
+        self.assertRaises(StopIteration, gen_orig.__next__)
+        gen_new = pickle.loads(p)
+        self.assertRaises(StopIteration, gen_new.__next__)
+
+
 if __name__ == '__main__':
     unittest.main()