GH-98363: Fix exception handling in batched() (GH-98523)

This commit is contained in:
Raymond Hettinger 2022-10-21 12:31:52 -05:00 committed by GitHub
parent ec1f6f5f13
commit a5ff80c8bc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 36 additions and 8 deletions

View file

@ -154,23 +154,36 @@ batched_next(batchedobject *bo)
if (result == NULL) {
return NULL;
}
iternextfunc iternext = *Py_TYPE(it)->tp_iternext;
PyObject **items = PySequence_Fast_ITEMS(result);
for (i=0 ; i < n ; i++) {
item = PyIter_Next(it);
item = iternext(it);
if (item == NULL) {
break;
goto null_item;
}
items[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;
}
PyList_SET_ITEM(result, i, item);
}
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 = {