gh-123471: make itertools.batched thread-safe (#129416)

This commit is contained in:
Pieter Eendebak 2025-03-12 11:00:33 +01:00 committed by GitHub
parent 155c44b901
commit 405a2d74cb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 51 additions and 2 deletions

View file

@ -191,12 +191,12 @@ batched_next(PyObject *op)
{
batchedobject *bo = batchedobject_CAST(op);
Py_ssize_t i;
Py_ssize_t n = bo->batch_size;
Py_ssize_t n = FT_ATOMIC_LOAD_SSIZE_RELAXED(bo->batch_size);
PyObject *it = bo->it;
PyObject *item;
PyObject *result;
if (it == NULL) {
if (n < 0) {
return NULL;
}
result = PyTuple_New(n);
@ -218,19 +218,28 @@ batched_next(PyObject *op)
if (PyErr_Occurred()) {
if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
/* Input raised an exception other than StopIteration */
FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
#ifndef Py_GIL_DISABLED
Py_CLEAR(bo->it);
#endif
Py_DECREF(result);
return NULL;
}
PyErr_Clear();
}
if (i == 0) {
FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
#ifndef Py_GIL_DISABLED
Py_CLEAR(bo->it);
#endif
Py_DECREF(result);
return NULL;
}
if (bo->strict) {
FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
#ifndef Py_GIL_DISABLED
Py_CLEAR(bo->it);
#endif
Py_DECREF(result);
PyErr_SetString(PyExc_ValueError, "batched(): incomplete batch");
return NULL;