From 5466d4af5fe76ec0a5fbc8a05675287d9e8e9d14 Mon Sep 17 00:00:00 2001 From: "T. Wouters" Date: Thu, 30 Mar 2017 09:58:35 -0700 Subject: [PATCH] bpo-29942: Fix the use of recursion in itertools.chain.from_iterable. (#889) 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. --- Lib/test/test_itertools.py | 8 ++++++ Misc/NEWS | 3 +++ Modules/itertoolsmodule.c | 52 ++++++++++++++++++++------------------ 3 files changed, 39 insertions(+), 24 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index ea1f57caad..c431f0dc6e 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1976,6 +1976,14 @@ class RegressionTests(unittest.TestCase): self.assertRaises(AssertionError, list, cycle(gen1())) self.assertEqual(hist, [0,1]) + def test_long_chain_of_empty_iterables(self): + # Make sure itertools.chain doesn't run into recursion limits when + # dealing with long chains of empty iterables. Even with a high + # number this would probably only fail in Py_DEBUG mode. + it = chain.from_iterable(() for unused in range(10000000)) + with self.assertRaises(StopIteration): + next(it) + class SubclassWithKwargsTest(unittest.TestCase): def test_keywords_in_subclass(self): # count is not subclassable... diff --git a/Misc/NEWS b/Misc/NEWS index 8646b8ce45..42a09b6d7b 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -301,6 +301,9 @@ Extension Modules Library ------- +- bpo-29942: Fix a crash in itertools.chain.from_iterable when encountering + long runs of empty iterables. + - bpo-10030: Sped up reading encrypted ZIP files by 2 times. - bpo-29204: Element.getiterator() and the html parameter of XMLParser() were diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index e91a0299c4..9823c77f77 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -1864,33 +1864,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 = (*Py_TYPE(lz->active)->tp_iternext)(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 = (*Py_TYPE(lz->active)->tp_iternext)(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; } static PyObject * -- 2.40.0