Issue #14288: Serialization support for builtin iterators.

1 files changed, 59 insertions, 1 deletions

diff --git a/Objects/iterobject.c b/Objects/iterobject.c
index 91a93f55be..bd0544ccd9 100644
--- a/Objects/iterobject.c
+++ b/Objects/iterobject.c
@@ -2,6 +2,19 @@
 
 #include "Python.h"
 
+/* Convenience function to get builtins.iter or builtins.reversed */
+PyObject *
+_PyIter_GetBuiltin(const char *iter)
+{
+    PyObject *mod, *attr;
+    mod = PyImport_ImportModule("builtins");
+    if (mod ==  NULL)
+        return NULL;
+    attr = PyObject_GetAttrString(mod, iter);
+    Py_DECREF(mod);
+    return attr;
+}
+
 typedef struct {
     PyObject_HEAD
     long      it_index;
@@ -88,8 +101,38 @@ iter_len(seqiterobject *it)
 
 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
 
+static PyObject *
+iter_reduce(seqiterobject *it)
+{
+    if (it->it_seq != NULL)
+        return Py_BuildValue("N(O)n", _PyIter_GetBuiltin("iter"),
+                             it->it_seq, it->it_index);
+    else
+        return Py_BuildValue("N(())", _PyIter_GetBuiltin("iter"));
+}
+
+PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
+
+static PyObject *
+iter_setstate(seqiterobject *it, PyObject *state)
+{
+    Py_ssize_t index = PyLong_AsSsize_t(state);
+    if (index == -1 && PyErr_Occurred())
+        return NULL;
+    if (it->it_seq != NULL) {
+        if (index < 0)
+            index = 0;
+        it->it_index = index;
+    }
+    Py_RETURN_NONE;
+}
+
+PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
+
 static PyMethodDef seqiter_methods[] = {
     {"__length_hint__", (PyCFunction)iter_len, METH_NOARGS, length_hint_doc},
+    {"__reduce__", (PyCFunction)iter_reduce, METH_NOARGS, reduce_doc},
+    {"__setstate__", (PyCFunction)iter_setstate, METH_O, setstate_doc},
     {NULL,              NULL}           /* sentinel */
 };
 
@@ -195,6 +238,21 @@ calliter_iternext(calliterobject *it)
     return NULL;
 }
 
+static PyObject *
+calliter_reduce(calliterobject *it)
+{
+    if (it->it_callable != NULL && it->it_sentinel != NULL)
+        return Py_BuildValue("N(OO)", _PyIter_GetBuiltin("iter"),
+                             it->it_callable, it->it_sentinel);
+    else
+        return Py_BuildValue("N(())", _PyIter_GetBuiltin("iter"));
+}
+
+static PyMethodDef calliter_methods[] = {
+    {"__reduce__", (PyCFunction)calliter_reduce, METH_NOARGS, reduce_doc},
+    {NULL,              NULL}           /* sentinel */
+};
+
 PyTypeObject PyCallIter_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     "callable_iterator",                        /* tp_name */
@@ -224,7 +282,7 @@ PyTypeObject PyCallIter_Type = {
     0,                                          /* tp_weaklistoffset */
     PyObject_SelfIter,                          /* tp_iter */
     (iternextfunc)calliter_iternext,            /* tp_iternext */
-    0,                                          /* tp_methods */
+    calliter_methods,                           /* tp_methods */
 };