Issue #14288: Serialization support for builtin iterators.

1 files changed, 53 insertions, 0 deletions

diff --git a/Objects/dictobject.c b/Objects/dictobject.c
index f0a07d7f3d..1924282b07 100644
--- a/Objects/dictobject.c
+++ b/Objects/dictobject.c
@@ -2323,9 +2323,16 @@ dictiter_len(dictiterobject *di)
 PyDoc_STRVAR(length_hint_doc,
              "Private method returning an estimate of len(list(it)).");
 
+static PyObject *
+dictiter_reduce(dictiterobject *di);
+
+PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
+
 static PyMethodDef dictiter_methods[] = {
     {"__length_hint__", (PyCFunction)dictiter_len, METH_NOARGS,
      length_hint_doc},
+     {"__reduce__", (PyCFunction)dictiter_reduce, METH_NOARGS,
+     reduce_doc},
     {NULL,              NULL}           /* sentinel */
 };
 
@@ -2560,6 +2567,52 @@ PyTypeObject PyDictIterItem_Type = {
 };
 
 
+static PyObject *
+dictiter_reduce(dictiterobject *di)
+{
+    PyObject *list;
+    dictiterobject tmp;
+
+    list = PyList_New(0);
+    if (!list)
+        return NULL;
+
+    /* copy the itertor state */
+    tmp = *di;
+    Py_XINCREF(tmp.di_dict);
+    
+    /* iterate the temporary into a list */
+    for(;;) {
+        PyObject *element = 0;
+        if (Py_TYPE(di) == &PyDictIterItem_Type)
+            element = dictiter_iternextitem(&tmp);
+        else if (Py_TYPE(di) == &PyDictIterKey_Type)
+            element = dictiter_iternextkey(&tmp);
+        else if (Py_TYPE(di) == &PyDictIterValue_Type)
+            element = dictiter_iternextvalue(&tmp);
+        else
+            assert(0);
+        if (element) {
+            if (PyList_Append(list, element)) {
+                Py_DECREF(element);
+                Py_DECREF(list);
+                Py_XDECREF(tmp.di_dict);
+                return NULL;
+            }
+            Py_DECREF(element);
+        } else
+            break;
+    }
+    Py_XDECREF(tmp.di_dict);
+    /* check for error */
+    if (tmp.di_dict != NULL) {
+        /* we have an error */
+        Py_DECREF(list);
+        return NULL;
+    }
+    return Py_BuildValue("N(N)", _PyIter_GetBuiltin("iter"), list);
+}
+
 /***********************************************/
 /* View objects for keys(), items(), values(). */
 /***********************************************/