8000 merge 3.3-slp (Stackless #130, fix exhausted generator pickling · akruis/cpython@895f2c6 · GitHub
[go: up one dir, main page]
Skip to content

Commit 895f2c6

Browse files
author
Anselm Kruis
committed
merge 3.3-slp (Stackless python#130, fix exhausted generator pickling
2 parents 9b28901 + 9fbbf8b commit 895f2c6

File tree

3 files changed

+69
-5
lines changed

3 files changed

+69
-5
lines changed

‎Stackless/changelog.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@ What's New in Stackless 3.X.X?
99

1010
*Release date: 20XX-XX-XX*
1111

12+
- https://bitbucket.org/stackless-dev/stackless/issues/130
13+
Pickling exhausted generators no longer crashes.
14+
1215

1316
What's New in Stackless 3.4.3?
1417
==============================

‎Stackless/pickling/prickelpit.c

Lines changed: 42 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1799,8 +1799,15 @@ static PyObject *
17991799
gen_reduce(PyGenObject *gen)
18001800
{
18011801
PyObject *tup;
1802-
PyObject *frame_reducer;
1803-
frame_reducer = slp_reduce_frame(gen->gi_frame);
1802+
PyObject *frame_reducer = (PyObject *)gen->gi_frame;
1803+
if (frame_reducer == NULL) {
1804+
/* Pickle NULL as None. See gen_setstate() for the corresponding
1805+
* unpickling code. */
1806+
Py_INCREF(Py_None);
1807+
frame_reducer = Py_None;
1808+
} else {
1809+
frame_reducer = slp_reduce_frame(gen->gi_frame);
1810+
}
18041811
if (frame_reducer == NULL)
18051812
return NULL;
18061813
tup = Py_BuildValue("(O()(Oi))",
@@ -1833,24 +1840,53 @@ gen_setstate(PyObject *self, PyObject *args)
18331840
{
18341841
PyGenObject *gen = (PyGenObject *) self;
18351842
PyFrameObject *f;
1843+
PyObject *obj;
18361844
int gi_running;
18371845

18381846
if (is_wrong_type(Py_TYPE(self))) return NULL;
18391847

18401848
if ((args = unwrap_frame_arg(args)) == NULL) /* now args is a counted ref! */
18411849
return NULL;
18421850

1843-
if (!PyArg_ParseTuple(args, "O!i:generator",
1844-
&PyFrame_Type, &f, &gi_running)) {
1851+
if (!PyArg_ParseTuple(args, "Oi:generator",
1852+
&obj, &gi_running)) {
18451853
Py_DECREF(args);
18461854
return NULL;
18471855
}
1848-
Py_DECREF(args);
1856+
1857+
if (obj == Py_None) {
1858+
/* No frame, generator is exhausted */
1859+
Py_CLEAR(gen->gi_frame);
1860+
1861+
/* Even if gi_frame is NULL, gi_code is still valid. Therefore
1862+
* I set it to the code of the exhausted frame singleton.
1863+
*/
1864+
assert(gen_exhausted_frame != NULL);
1865+
assert(PyFrame_Check(gen_exhausted_frame));
1866+
obj = (PyObject *)gen_exhausted_frame->f_code;
1867+
Py_INCREF(obj);
1868+
Py_SETREF(gen->gi_code, obj);
1869+
1870+
gen->gi_running = gi_running;
1871+
Py_TYPE(gen) = Py_TYPE(gen)->tp_base;
1872+
Py_INCREF(gen);
1873+
Py_DECREF(args); /* holds the obj ref */
1874+
return (PyObject *)gen;
1875+
}
1876+
if (!PyFrame_Check(obj)) {
1877+
PyErr_SetString(PyExc_TypeError,
1878+
"invalid frame object for gen_setstate");
1879+
Py_DECREF(args); /* holds the obj ref */
1880+
return NULL;
1881+
}
1882+
f = (PyFrameObject *)obj;
1883+
obj = NULL;
18491884

18501885
if (!gi_running) {
18511886
if ((f = slp_ensure_new_frame(f)) != NULL) {
18521887
/* use a second one for late initialization */
18531888
PyGenObject *tmpgen;
1889+
Py_DECREF(args); /* holds the frame ref */
18541890
/* PyGenerator_New eats an existing reference */
18551891
if ((tmpgen = (PyGenObject *)
18561892
PyGen_New(f)) == NULL) {
@@ -1884,6 +1920,7 @@ gen_setstate(PyObject *self, PyObject *args)
18841920
* generator without the corresponding tasklet.
18851921
*/
18861922
Py_INCREF(f);
1923+
Py_DECREF(args); /* holds the frame ref */
18871924
Py_SETREF(gen->gi_frame, f);
18881925
Py_INCREF(f->f_code);
18891926
Py_SETREF(gen->gi_code, (PyObject *)f->f_code);

‎Stackless/unittests/test_generator.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import gc
44
import stackless
55
import types
6+
import pickle
67

78
from support import test_main # @UnusedImport
89
from support import StacklessTestCase
@@ -51,5 +52,28 @@ def test_wrap_generator_frame_code(self):
5152
self.assertIsNot(g0.gi_frame, g1.gi_frame)
5253
self.assertEqual(g0.__name__, "exhausted_generator")
5354

55+
56+
class TestGeneratorPickling(StacklessTestCase):
57+
def test_name(self):
58+
def g():
59+
yield 1
60+
gen_orig = g()
61+
p = pickle.dumps(gen_orig)
62+
gen_new = pickle.loads(p)
63+
64+
self.assertEqual(gen_new.__name__, gen_orig.__name__)
65+
66+
def test_exhausted(self):
67+
def g():
68+
yield 1
69+
gen_orig = g()
70+
self.assertEqual(gen_orig.__next__(), 1)
71+
self.assertRaises(StopIteration, gen_orig.__next__)
72+
p = pickle.dumps(gen_orig)
73+
self.assertRaises(StopIteration, gen_orig.__next__)
74+
gen_new = pickle.loads(p)
75+
self.assertRaises(StopIteration, gen_new.__next__)
76+
77+
5478
if __name__ == '__main__':
5579
unittest.main()

0 commit comments

Comments
 (0)
0