bpo-34574: Prevent OrderedDict iterators from exhaustion during pickling. (GH-9051)

This commit is contained in:
Sergey Fedoseev 2018-10-20 10:20:39 +05:00 committed by Serhiy Storchaka
parent 8c9fd9c91b
commit a5259fb05d
3 changed files with 28 additions and 28 deletions

View File

@ -732,6 +732,23 @@ class CPythonOrderedDictTests(OrderedDictTests, unittest.TestCase):
del od['c']
self.assertEqual(list(od), list('bdeaf'))
def test_iterators_pickling(self):
OrderedDict = self.OrderedDict
pairs = [('c', 1), ('b', 2), ('a', 3), ('d', 4), ('e', 5), ('f', 6)]
od = OrderedDict(pairs)
for method_name in ('keys', 'values', 'items'):
meth = getattr(od, method_name)
expected = list(meth())[1:]
for i in range(pickle.HIGHEST_PROTOCOL + 1):
with self.subTest(method_name=method_name, protocol=i):
it = iter(meth())
next(it)
p = pickle.dumps(it, i)
unpickled = pickle.loads(p)
self.assertEqual(list(unpickled), expected)
self.assertEqual(list(it), expected)
class PurePythonOrderedDictSubclassTests(PurePythonOrderedDictTests):

View File

@ -0,0 +1,2 @@
OrderedDict iterators are not exhausted during pickling anymore. Patch by
Sergey Fedoseev.

View File

@ -1805,38 +1805,19 @@ PyDoc_STRVAR(reduce_doc, "Return state information for pickling");
static PyObject *
odictiter_reduce(odictiterobject *di)
{
PyObject *list, *iter;
list = PyList_New(0);
if (!list)
return NULL;
/* copy the iterator state */
odictiterobject tmp = *di;
Py_XINCREF(tmp.di_odict);
Py_XINCREF(tmp.di_current);
/* iterate the temporary into a list */
for(;;) {
PyObject *element = odictiter_iternext(di);
if (element) {
if (PyList_Append(list, element)) {
Py_DECREF(element);
Py_DECREF(list);
return NULL;
}
Py_DECREF(element);
}
else {
/* done iterating? */
break;
}
}
if (PyErr_Occurred()) {
Py_DECREF(list);
PyObject *list = PySequence_List((PyObject*)&tmp);
Py_XDECREF(tmp.di_odict);
Py_XDECREF(tmp.di_current);
if (list == NULL) {
return NULL;
}
iter = _PyObject_GetBuiltin("iter");
if (iter == NULL) {
Py_DECREF(list);
return NULL;
}
return Py_BuildValue("N(N)", iter, list);
return Py_BuildValue("N(N)", _PyObject_GetBuiltin("iter"), list);
}
static PyMethodDef odictiter_methods[] = {