Make itertools.batched thread-safe

Author: eendebakpt

Modules/itertoolsmodule.c

@@ -186,12 +186,12 @@ static PyObject *
 batched_next(batchedobject *bo)
 {
     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);
@@ -213,19 +213,19 @@ batched_next(batchedobject *bo)
     if (PyErr_Occurred()) {
         if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
             /* Input raised an exception other than StopIteration */
-            Py_CLEAR(bo->it);
+            FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
             Py_DECREF(result);
             return NULL;
         }
         PyErr_Clear();
     }
     if (i == 0) {
-        Py_CLEAR(bo->it);
+        FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
         Py_DECREF(result);
         return NULL;
     }
     if (bo->strict) {
-        Py_CLEAR(bo->it);
+        FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
         Py_DECREF(result);
         PyErr_SetString(PyExc_ValueError, "batched(): incomplete batch");
         return NULL;