From ccbb4021b4543850a342f5b032ee8be0e7970e93 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Fri, 21 Oct 2022 10:26:02 -0500 Subject: [PATCH 1/5] Add failing test for unhandled iterator expection --- Lib/test/test_itertools.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index c0e35711a2b3dd..a0a740fba8e8e3 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -2012,6 +2012,20 @@ def __iter__(self): def __next__(self): 3 // 0 +class E2: + 'Test propagation of exceptions after two iterations' + def __init__(self, seqn): + self.seqn = seqn + self.i = 0 + def __iter__(self): + return self + def __next__(self): + if self.i == 2: + raise ZeroDivisionError + v = self.seqn[self.i] + self.i += 1 + return v + class S: 'Test immediate stop' def __init__(self, seqn): @@ -2050,6 +2064,7 @@ def test_batched(self): self.assertRaises(TypeError, batched, X(s), 2) self.assertRaises(TypeError, batched, N(s), 2) self.assertRaises(ZeroDivisionError, list, batched(E(s), 2)) + self.assertRaises(ZeroDivisionError, list, batched(E2(s), 4)) def test_chain(self): for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)): From 9343e27f42842b40bc1cf2372f5fda8f705ba5b5 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Fri, 21 Oct 2022 10:32:21 -0500 Subject: [PATCH 2/5] Fix failing test. --- Modules/itertoolsmodule.c | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 868e8a8b384f1b..b9ee0982f0dacc 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -155,8 +155,18 @@ batched_next(batchedobject *bo) return NULL; } for (i=0 ; i < n ; i++) { - item = PyIter_Next(it); + item = (*Py_TYPE(it)->tp_iternext)(it); if (item == NULL) { + if (PyErr_Occurred()) { + if (PyErr_ExceptionMatches(PyExc_StopIteration)) { + PyErr_Clear(); + } else { + /* input raised an exception other than StopIteration */ + Py_CLEAR(bo->it); + Py_DECREF(result); + return NULL; + } + } break; } PyList_SET_ITEM(result, i, item); From 9d8cf66922b978adf0519d18bce5666e09512891 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Fri, 21 Oct 2022 10:42:07 -0500 Subject: [PATCH 3/5] Refactor so that the logic is more clear --- Modules/itertoolsmodule.c | 33 +++++++++++++++++---------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index b9ee0982f0dacc..cd221ff48cc636 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -157,30 +157,31 @@ batched_next(batchedobject *bo) for (i=0 ; i < n ; i++) { item = (*Py_TYPE(it)->tp_iternext)(it); if (item == NULL) { - if (PyErr_Occurred()) { - if (PyErr_ExceptionMatches(PyExc_StopIteration)) { - PyErr_Clear(); - } else { - /* input raised an exception other than StopIteration */ - Py_CLEAR(bo->it); - Py_DECREF(result); - return NULL; - } - } - break; + goto null_item; } PyList_SET_ITEM(result, i, item); } + return result; + + null_item: + if (PyErr_Occurred()) { + if (PyErr_ExceptionMatches(PyExc_StopIteration)) { + PyErr_Clear(); + } else { + /* input raised an exception other than StopIteration */ + Py_CLEAR(bo->it); + Py_DECREF(result); + return NULL; + } + } if (i == 0) { Py_CLEAR(bo->it); Py_DECREF(result); return NULL; } - if (i < n) { - PyObject *short_list = PyList_GetSlice(result, 0, i); - Py_SETREF(result, short_list); - } - return result; + PyObject *short_list = PyList_GetSlice(result, 0, i); + Py_DECREF(result); + return short_list; } static PyTypeObject batched_type = { From 699bdf679228c2595a40245db69d3498f503b8df Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Fri, 21 Oct 2022 10:56:01 -0500 Subject: [PATCH 4/5] Loop invariant code motion: hoist the slot lookup out of the loop. --- Modules/itertoolsmodule.c | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index cd221ff48cc636..1e5f9d142810b5 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -154,8 +154,9 @@ batched_next(batchedobject *bo) if (result == NULL) { return NULL; } + iternextfunc iternext = *Py_TYPE(it)->tp_iternext; for (i=0 ; i < n ; i++) { - item = (*Py_TYPE(it)->tp_iternext)(it); + item = iternext(it); if (item == NULL) { goto null_item; } From 4b74a0b1b38d076642b90dc83f073d1a1bf3cda6 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Fri, 21 Oct 2022 11:55:42 -0500 Subject: [PATCH 5/5] Loop invariant code motion: hoist ob_item lookup out of the loop. --- Modules/itertoolsmodule.c | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 1e5f9d142810b5..627e698fc6b9d8 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -155,12 +155,13 @@ batched_next(batchedobject *bo) return NULL; } iternextfunc iternext = *Py_TYPE(it)->tp_iternext; + PyObject **items = PySequence_Fast_ITEMS(result); for (i=0 ; i < n ; i++) { item = iternext(it); if (item == NULL) { goto null_item; } - PyList_SET_ITEM(result, i, item); + items[i] = item; } return result;