From a5259fb05d03f4871837c14fed704541a20896c0 Mon Sep 17 00:00:00 2001
From: Sergey Fedoseev <fedoseev.sergey@gmail.com>
Date: Sat, 20 Oct 2018 10:20:39 +0500
Subject: [PATCH] bpo-34574: Prevent OrderedDict iterators from exhaustion
 during pickling. (GH-9051)

---
 Lib/test/test_ordered_dict.py                 | 17 +++++++++
 .../2018-09-04-09-32-54.bpo-34574.X4RwYI.rst  |  2 +
 Objects/odictobject.c                         | 37 +++++--------------
 3 files changed, 28 insertions(+), 28 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2018-09-04-09-32-54.bpo-34574.X4RwYI.rst

diff --git a/Lib/test/test_ordered_dict.py b/Lib/test/test_ordered_dict.py
index a296f6055b1a12..b1d7f86a6760b3 100644
--- a/Lib/test/test_ordered_dict.py
+++ b/Lib/test/test_ordered_dict.py
@@ -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):
 
diff --git a/Misc/NEWS.d/next/Library/2018-09-04-09-32-54.bpo-34574.X4RwYI.rst b/Misc/NEWS.d/next/Library/2018-09-04-09-32-54.bpo-34574.X4RwYI.rst
new file mode 100644
index 00000000000000..de718ad1e6f102
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2018-09-04-09-32-54.bpo-34574.X4RwYI.rst
@@ -0,0 +1,2 @@
+OrderedDict iterators are not exhausted during pickling anymore. Patch by
+Sergey Fedoseev.
diff --git a/Objects/odictobject.c b/Objects/odictobject.c
index 47e77b6046b24f..67c3674ee09528 100644
--- a/Objects/odictobject.c
+++ b/Objects/odictobject.c
@@ -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[] = {