8000 bpo-34574: Prevent OrderedDict iterators from exhaustion during pickl… · python/cpython@0d3dd9f · GitHub
[go: up one dir, main page]

Skip to content

Commit 0d3dd9f

Browse files
bpo-34574: Prevent OrderedDict iterators from exhaustion during pickling. (GH-9051)
(cherry picked from commit a5259fb) Co-authored-by: Sergey Fedoseev <fedoseev.sergey@gmail.com>
1 parent 42c52a9 commit 0d3dd9f

File tree

3 files changed

+28
-28
lines changed

3 files changed

+28
-28
lines changed

Lib/test/test_ordered_dict.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -732,6 +732,23 @@ def test_key_change_during_iteration(self):
732732
del od['c']
733733
self.assertEqual(list(od), list('bdeaf'))
734734

735+
def test_iterators_pickling(self):
736+
OrderedDict = self.OrderedDict
737+
pairs = [('c', 1), ('b', 2), ('a', 3), ('d', 4), ('e', 5), ('f', 6)]
738+
od = OrderedDict(pairs)
739+
740+
for method_name in ('keys', 'values', 'items'):
741+
meth = getattr(od, method_name)
742+
expected = list(meth())[1:]
743+
for i in range(pickle.HIGHEST_PROTOCOL + 1):
744+
with self.subTest(method_name=method_name, protocol=i):
745+
it = iter(meth())
746+
next(it)
747+
p = pickle.dumps(it, i)
748+
unpickled = pickle.loads(p)
749+
self.assertEqual(list(unpickled), expected)
750+
self.assertEqual(list(it), expected)
751+
735752

736753
class PurePythonOrderedDictSubclassTests(PurePythonOrderedDictTests):
737754

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
OrderedDict iterators are not exhausted during pickling anymore. Patch by
2+
Sergey Fedoseev.

Objects/odictobject.c

Lines changed: 9 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1923,38 +1923,19 @@ PyDoc_STRVAR(reduce_doc, "Return state information for pickling");
19231923
static PyObject *
19241924
odictiter_reduce(odictiterobject *di)
19251925
{
1926-
PyObject *list, *iter;
1927-
1928-
list = PyList_New(0);
1929-
if (!list)
1930-
return NULL;
1926+
/* copy the iterator state */
1927+
odictiterobject tmp = *di;
1928+
Py_XINCREF(tmp.di_odict);
1929+
Py_XINCREF(tmp.di_current);
19311930

19321931
/* iterate the temporary into a list */
1933-
for(;;) {
1934-
PyObject *element = odictiter_iternext(di);
1935-
if (element) {
1936-
if (PyList_Append(list, element)) {
1937-
Py_DECREF(element);
1938-
Py_DECREF(list);
1939-
return NULL;
1940-
}
1941-
Py_DECREF(element);
1942-
}
1943-
else {
1944-
/* done iterating? */
1945-
break;
1946-
}
1947-
}
1948-
if (PyErr_Occurred()) {
1949-
Py_DECREF(list);
1950-
return NULL;
1951-
}
1952-
iter = _PyObject_GetBuiltin("iter");
1953-
if (iter == NULL) {
1954-
Py_DECREF(list);
1932+
PyObject *list = PySequence_List((PyObject*)&tmp);
1933+
Py_XDECREF(tmp.di_odict);
1934+
Py_XDECREF(tmp.di_current);
1935+
if (list == NULL) {
19551936
return NULL;
19561937
}
1957-
return Py_BuildValue("N(N)", iter, list);
1938+
return Py_BuildValue("N(N)", _PyObject_GetBuiltin("iter"), list);
19581939
}
19591940

19601941
static PyMethodDef odictiter_methods[] = {

0 commit comments

Comments
 (0)
0