Skip to content

Commit a5ff80c

Browse files
authored
GH-98363: Fix exception handling in batched() (GH-98523)
1 parent ec1f6f5 commit a5ff80c

File tree

2 files changed

+36
-8
lines changed

2 files changed

+36
-8
lines changed

Lib/test/test_itertools.py

+15
Original file line numberDiff line numberDiff line change
@@ -2012,6 +2012,20 @@ def __iter__(self):
20122012
def __next__(self):
20132013
3 // 0
20142014

2015+
class E2:
2016+
'Test propagation of exceptions after two iterations'
2017+
def __init__(self, seqn):
2018+
self.seqn = seqn
2019+
self.i = 0
2020+
def __iter__(self):
2021+
return self
2022+
def __next__(self):
2023+
if self.i == 2:
2024+
raise ZeroDivisionError
2025+
v = self.seqn[self.i]
2026+
self.i += 1
2027+
return v
2028+
20152029
class S:
20162030
'Test immediate stop'
20172031
def __init__(self, seqn):
@@ -2050,6 +2064,7 @@ def test_batched(self):
20502064
self.assertRaises(TypeError, batched, X(s), 2)
20512065
self.assertRaises(TypeError, batched, N(s), 2)
20522066
self.assertRaises(ZeroDivisionError, list, batched(E(s), 2))
2067+
self.assertRaises(ZeroDivisionError, list, batched(E2(s), 4))
20532068

20542069
def test_chain(self):
20552070
for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):

Modules/itertoolsmodule.c

+21-8
Original file line numberDiff line numberDiff line change
@@ -154,23 +154,36 @@ batched_next(batchedobject *bo)
154154
if (result == NULL) {
155155
return NULL;
156156
}
157+
iternextfunc iternext = *Py_TYPE(it)->tp_iternext;
158+
PyObject **items = PySequence_Fast_ITEMS(result);
157159
for (i=0 ; i < n ; i++) {
158-
item = PyIter_Next(it);
160+
item = iternext(it);
159161
if (item == NULL) {
160-
break;
162+
goto null_item;
163+
}
164+
items[i] = item;
165+
}
166+
return result;
167+
168+
null_item:
169+
if (PyErr_Occurred()) {
170+
if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
171+
PyErr_Clear();
172+
} else {
173+
/* input raised an exception other than StopIteration */
174+
Py_CLEAR(bo->it);
175+
Py_DECREF(result);
176+
return NULL;
161177
}
162-
PyList_SET_ITEM(result, i, item);
163178
}
164179
if (i == 0) {
165180
Py_CLEAR(bo->it);
166181
Py_DECREF(result);
167182
return NULL;
168183
}
169-
if (i < n) {
170-
PyObject *short_list = PyList_GetSlice(result, 0, i);
171-
Py_SETREF(result, short_list);
172-
}
173-
return result;
184+
PyObject *short_list = PyList_GetSlice(result, 0, i);
185+
Py_DECREF(result);
186+
return short_list;
174187
}
175188

176189
static PyTypeObject batched_type = {

0 commit comments

Comments
 (0)