SF #950057: itertools.chain doesn't "process" exceptions as they occur
Both cycle() and chain() were handling exceptions only when switching
input sources. The patch makes the handle more immediate.
Will backport.
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 73e8809..54e46e1 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -644,6 +644,36 @@
self.assertEqual(first, second)
+ def test_sf_950057(self):
+ # Make sure that chain() and cycle() catch exceptions immediately
+ # rather than when shifting between input sources
+
+ def gen1():
+ hist.append(0)
+ yield 1
+ hist.append(1)
+ assert False
+ hist.append(2)
+
+ def gen2(x):
+ hist.append(3)
+ yield 2
+ hist.append(4)
+ if x:
+ raise StopIteration
+
+ hist = []
+ self.assertRaises(AssertionError, list, chain(gen1(), gen2(False)))
+ self.assertEqual(hist, [0,1])
+
+ hist = []
+ self.assertRaises(AssertionError, list, chain(gen1(), gen2(True)))
+ self.assertEqual(hist, [0,1])
+
+ hist = []
+ self.assertRaises(AssertionError, list, cycle(gen1()))
+ self.assertEqual(hist, [0,1])
+
libreftest = """ Doctest for examples in the library reference: libitertools.tex
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 4ce4643..3515bc6 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -699,6 +699,12 @@
PyList_Append(lz->saved, item);
return item;
}
+ if (PyErr_Occurred()) {
+ if (PyErr_ExceptionMatches(PyExc_StopIteration))
+ PyErr_Clear();
+ else
+ return NULL;
+ }
if (PyList_Size(lz->saved) == 0)
return NULL;
it = PyObject_GetIter(lz->saved);
@@ -1658,6 +1664,12 @@
item = PyIter_Next(it);
if (item != NULL)
return item;
+ if (PyErr_Occurred()) {
+ if (PyErr_ExceptionMatches(PyExc_StopIteration))
+ PyErr_Clear();
+ else
+ return NULL;
+ }
lz->iternum++;
}
return NULL;