diff options
Diffstat (limited to 'Modules/itertoolsmodule.c')
| -rw-r--r-- | Modules/itertoolsmodule.c | 13 |
1 files changed, 11 insertions, 2 deletions
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 3e425ee..40e436a6 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -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; |