Skip to content

Commit d694a06

Browse files
authored
bpo-29942: Fix the use of recursion in itertools.chain.from_iterable. (#913)
Fix the use of recursion in itertools.chain.from_iterable. Using recursion is unnecessary, and can easily cause stack overflows, especially when building in low optimization modes or with Py_DEBUG enabled. (cherry picked from commit 5466d4a)
1 parent 079f21f commit d694a06

File tree

3 files changed

+39
-24
lines changed

3 files changed

+39
-24
lines changed

Lib/test/test_itertools.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1465,6 +1465,14 @@ def gen2(x):
14651465
self.assertRaises(AssertionError, list, cycle(gen1()))
14661466
self.assertEqual(hist, [0,1])
14671467

1468+
def test_long_chain_of_empty_iterables(self):
1469+
# Make sure itertools.chain doesn't run into recursion limits when
1470+
# dealing with long chains of empty iterables. Even with a high
1471+
# number this would probably only fail in Py_DEBUG mode.
1472+
it = chain.from_iterable(() for unused in xrange(10000000))
1473+
with self.assertRaises(StopIteration):
1474+
next(it)
1475+
14681476
class SubclassWithKwargsTest(unittest.TestCase):
14691477
def test_keywords_in_subclass(self):
14701478
# count is not subclassable...

Misc/NEWS

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,9 @@ Extension Modules
4242
Library
4343
-------
4444

45+
- bpo-29942: Fix a crash in itertools.chain.from_iterable when encountering
46+
long runs of empty iterables.
47+
4548
- bpo-29861: Release references to tasks, their arguments and their results
4649
as soon as they are finished in multiprocessing.Pool.
4750

Modules/itertoolsmodule.c

Lines changed: 28 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1708,33 +1708,37 @@ chain_next(chainobject *lz)
17081708
{
17091709
PyObject *item;
17101710

1711-
if (lz->source == NULL)
1712-
return NULL; /* already stopped */
1713-
1714-
if (lz->active == NULL) {
1715-
PyObject *iterable = PyIter_Next(lz->source);
1716-
if (iterable == NULL) {
1717-
Py_CLEAR(lz->source);
1718-
return NULL; /* no more input sources */
1719-
}
1720-
lz->active = PyObject_GetIter(iterable);
1721-
Py_DECREF(iterable);
1711+
/* lz->source is the iterator of iterables. If it's NULL, we've already
1712+
* consumed them all. lz->active is the current iterator. If it's NULL,
1713+
* we should grab a new one from lz->source. */
1714+
while (lz->source != NULL) {
17221715
if (lz->active == NULL) {
1723-
Py_CLEAR(lz->source);
1724-
return NULL; /* input not iterable */
1716+
PyObject *iterable = PyIter_Next(lz->source);
1717+
if (iterable == NULL) {
1718+
Py_CLEAR(lz->source);
1719+
return NULL; /* no more input sources */
1720+
}
1721+
lz->active = PyObject_GetIter(iterable);
1722+
Py_DECREF(iterable);
1723+
if (lz->active == NULL) {
1724+
Py_CLEAR(lz->source);
1725+
return NULL; /* input not iterable */
1726+
}
17251727
}
1728+
item = PyIter_Next(lz->active);
1729+
if (item != NULL)
1730+
return item;
1731+
if (PyErr_Occurred()) {
1732+
if (PyErr_ExceptionMatches(PyExc_StopIteration))
1733+
PyErr_Clear();
1734+
else
1735+
return NULL; /* input raised an exception */
1736+
}
1737+
/* lz->active is consumed, try with the next iterable. */
1738+
Py_CLEAR(lz->active);
17261739
}
1727-
item = PyIter_Next(lz->active);
1728-
if (item != NULL)
1729-
return item;
1730-
if (PyErr_Occurred()) {
1731-
if (PyErr_ExceptionMatches(PyExc_StopIteration))
1732-
PyErr_Clear();
1733-
else
1734-
return NULL; /* input raised an exception */
1735-
}
1736-
Py_CLEAR(lz->active);
1737-
return chain_next(lz); /* recurse and use next active */
1740+
/* Everything had been consumed already. */
1741+
return NULL;
17381742
}
17391743

17401744
PyDoc_STRVAR(chain_doc,

0 commit comments

Comments
 (0)