From d694a06206fc09b76b4507aacde5e69a248f434f Mon Sep 17 00:00:00 2001 From: "T. Wouters" Date: Thu, 30 Mar 2017 12:49:22 -0700 Subject: [PATCH] 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 5466d4af5fe76ec0a5fbc8a05675287d9e8e9d14) --- 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 5672104257..7b2362d29e 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1465,6 +1465,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 xrange(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 0ca08d6bef..559a8088a7 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -42,6 +42,9 @@ Extension Modules Library ------- +- bpo-29942: Fix a crash in itertools.chain.from_iterable when encountering + long runs of empty iterables. + - bpo-29861: Release references to tasks, their arguments and their results as soon as they are finished in multiprocessing.Pool. 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, -- 2.50.1