diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index cf579d4da4e0df..677b66fe703e1c 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -157,6 +157,31 @@ def test_accumulate(self): with self.assertRaises(TypeError): list(accumulate([10, 20], 100)) + def test_accumulate_reenter(self): + class I: + count = 0 + def __iter__(self): + return self + def __next__(self): + self.count += 1 + if self.count == 2: + return next(it) + return self.count + + it = accumulate(I()) + self.assertEqual(next(it), 1) + with self.assertRaisesRegex(RuntimeError, "accumulate"): + next(it) + + def test_accumulate_reenter_func(self): + def reenter(a, b): + return next(it) + + it = accumulate([1, 2, 3], reenter) + self.assertEqual(next(it), 1) + with self.assertRaisesRegex(RuntimeError, "accumulate"): + next(it) + def test_batched(self): self.assertEqual(list(batched('ABCDEFG', 3)), [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]) diff --git a/Misc/NEWS.d/next/Library/2026-07-23-22-00-00.gh-issue-154570.accumulate-reentrancy.rst b/Misc/NEWS.d/next/Library/2026-07-23-22-00-00.gh-issue-154570.accumulate-reentrancy.rst new file mode 100644 index 00000000000000..116dc6d2d7083a --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-23-22-00-00.gh-issue-154570.accumulate-reentrancy.rst @@ -0,0 +1,2 @@ +Fix :func:`itertools.accumulate` silently corrupting its running total +on reentrancy; it now raises :exc:`RuntimeError` instead. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index c0023c839ca7fe..3f95761df377cc 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3079,6 +3079,7 @@ typedef struct { PyObject *binop; PyObject *initial; itertools_state *state; + uint8_t running; } accumulateobject; #define accumulateobject_CAST(op) ((accumulateobject *)(op)) @@ -3120,6 +3121,7 @@ itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable, lz->it = it; lz->initial = Py_XNewRef(initial); lz->state = find_state_by_type(type); + lz->running = 0; return (PyObject *)lz; } @@ -3160,12 +3162,21 @@ accumulate_next_lock_held(PyObject *op) lz->initial = Py_NewRef(Py_None); return Py_NewRef(lz->total); } + if (lz->running) { + PyErr_SetString(PyExc_RuntimeError, + "cannot re-enter the accumulate iterator"); + return NULL; + } + lz->running = 1; val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it); - if (val == NULL) + if (val == NULL) { + lz->running = 0; return NULL; + } if (lz->total == NULL) { lz->total = Py_NewRef(val); + lz->running = 0; return lz->total; } @@ -3174,6 +3185,7 @@ accumulate_next_lock_held(PyObject *op) else newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL); Py_DECREF(val); + lz->running = 0; if (newtotal == NULL) return NULL;