diff options
Diffstat (limited to 'Modules')
| -rw-r--r-- | Modules/itertoolsmodule.c | 52 |
1 files changed, 28 insertions, 24 deletions
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index f91501bb24..cf3aadfd3e 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -1708,33 +1708,37 @@ chain_next(chainobject *lz) { PyObject *item; - if (lz->source == NULL) - return NULL; /* already stopped */ - - if (lz->active == NULL) { - PyObject *iterable = PyIter_Next(lz->source); - if (iterable == NULL) { - Py_CLEAR(lz->source); - return NULL; /* no more input sources */ - } - lz->active = PyObject_GetIter(iterable); - Py_DECREF(iterable); + /* lz->source is the iterator of iterables. If it's NULL, we've already + * consumed them all. lz->active is the current iterator. If it's NULL, + * we should grab a new one from lz->source. */ + while (lz->source != NULL) { if (lz->active == NULL) { - Py_CLEAR(lz->source); - return NULL; /* input not iterable */ + PyObject *iterable = PyIter_Next(lz->source); + if (iterable == NULL) { + Py_CLEAR(lz->source); + return NULL; /* no more input sources */ + } + lz->active = PyObject_GetIter(iterable); + Py_DECREF(iterable); + if (lz->active == NULL) { + Py_CLEAR(lz->source); + return NULL; /* input not iterable */ + } } + item = PyIter_Next(lz->active); + if (item != NULL) + return item; + if (PyErr_Occurred()) { + if (PyErr_ExceptionMatches(PyExc_StopIteration)) + PyErr_Clear(); + else + return NULL; /* input raised an exception */ + } + /* lz->active is consumed, try with the next iterable. */ + Py_CLEAR(lz->active); } - item = PyIter_Next(lz->active); - if (item != NULL) - return item; - if (PyErr_Occurred()) { - if (PyErr_ExceptionMatches(PyExc_StopIteration)) - PyErr_Clear(); - else - return NULL; /* input raised an exception */ - } - Py_CLEAR(lz->active); - return chain_next(lz); /* recurse and use next active */ + /* Everything had been consumed already. */ + return NULL; } PyDoc_STRVAR(chain_doc, |