diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -861,9 +861,10 @@ typedef struct { PyObject_HEAD + int firstpass; + Py_ssize_t index; + PyObject *saved; PyObject *it; - PyObject *saved; - int firstpass; } cycleobject; static PyTypeObject cycle_type; @@ -903,6 +904,7 @@ lz->it = it; lz->saved = saved; lz->firstpass = 0; + lz->index = 0; return (PyObject *)lz; } @@ -919,7 +921,8 @@ static int cycle_traverse(cycleobject *lz, visitproc visit, void *arg) { - Py_VISIT(lz->it); + if (lz->it) + Py_VISIT(lz->it); Py_VISIT(lz->saved); return 0; } @@ -928,13 +931,11 @@ cycle_next(cycleobject *lz) { PyObject *item; - PyObject *it; - PyObject *tmp; - - while (1) { + + if (!lz->firstpass) { item = PyIter_Next(lz->it); if (item != NULL) { - if (!lz->firstpass && PyList_Append(lz->saved, item)) { + if (PyList_Append(lz->saved, item)) { Py_DECREF(item); return NULL; } @@ -942,27 +943,47 @@ } /* Note: StopIteration is already cleared by PyIter_Next() */ if (PyErr_Occurred()) - return NULL; - if (PyList_Size(lz->saved) == 0) return NULL; - it = PyObject_GetIter(lz->saved); - if (it == NULL) - return NULL; - tmp = lz->it; - lz->it = it; lz->firstpass = 1; - Py_DECREF(tmp); + Py_CLEAR(lz->it); } + if (Py_SIZE(lz->saved) == 0) + return NULL; + if (lz->index >= Py_SIZE(lz->saved)) + lz->index = 0; + item = PyList_GET_ITEM(lz->saved, lz->index); + lz->index++; + Py_INCREF(item); + return item; } static PyObject * cycle_reduce(cycleobject *lz) { - /* Create a new cycle with the iterator tuple, then set - * the saved state on it. - */ - return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), - lz->it, lz->saved, lz->firstpass); + /* Create a new cycle with the iterator tuple, then set the saved state */ + if (lz->it == NULL) { + PyObject *seq, *it, *result; + + if (lz->index >= Py_SIZE(lz->saved)) + lz->index = 0; + seq = PyList_GetSlice(lz->saved, lz->index, Py_SIZE(lz->saved)); + if (seq == NULL) + return NULL; + it = PyObject_GetIter(seq); + Py_DECREF(seq); + if (it == NULL) + return NULL; + seq = PyList_GetSlice(lz->saved, 0, lz->index); + if (seq == NULL) { + Py_DECREF(it); + return NULL; + } + result = Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), it, seq, 0); + Py_DECREF(it); + Py_DECREF(seq); + return result; + } + return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved, 0); } static PyObject *