Skip to content

Commit

Permalink
bpo-34574: Prevent OrderedDict iterators from exhaustion during pickl…
Browse files Browse the repository at this point in the history
…ing. (pythonGH-9051)
  • Loading branch information
sir-sigurd authored and serhiy-storchaka committed Oct 20, 2018
1 parent 8c9fd9c commit a5259fb
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 28 deletions.
17 changes: 17 additions & 0 deletions Lib/test/test_ordered_dict.py
Original file line number Diff line number Diff line change
Expand Up @@ -732,6 +732,23 @@ def test_key_change_during_iteration(self):
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):

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
OrderedDict iterators are not exhausted during pickling anymore. Patch by
Sergey Fedoseev.
37 changes: 9 additions & 28 deletions Objects/odictobject.c
Original file line number Diff line number Diff line change
Expand Up @@ -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);
return NULL;
}
iter = _PyObject_GetBuiltin("iter");
if (iter == NULL) {
Py_DECREF(list);
PyObject *list = PySequence_List((PyObject*)&tmp);
Py_XDECREF(tmp.di_odict);
Py_XDECREF(tmp.di_current);
if (list == NULL) {
return NULL;
}
return Py_BuildValue("N(N)", iter, list);
return Py_BuildValue("N(N)", _PyObject_GetBuiltin("iter"), list);
}

static PyMethodDef odictiter_methods[] = {
Expand Down

0 comments on commit a5259fb

Please sign in to comment.