Skip to content

Commit 643254e

Browse files
committed
gh-154570: Fix itertools.accumulate() silently corrupting its total on reentrancy
A source iterable or func= callable that calls back into next() on the same accumulate object mid-step silently corrupted the running total instead of erroring. Adds a running guard field to accumulateobject, mirroring the existing pattern already used by teedataobject for the identical class of bug (see teedataobject_getitem_lock_held). accumulate now raises RuntimeError("cannot re-enter the accumulate iterator") on reentrant access instead of silently corrupting state, matching tee's existing behavior for the same defect class. Adds a regression test (test_accumulate_reenter) in the same style as the existing test_tee_reenter.
1 parent 66e313f commit 643254e

3 files changed

Lines changed: 38 additions & 1 deletion

File tree

Lib/test/test_itertools.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,25 @@ def test_accumulate(self):
157157
with self.assertRaises(TypeError):
158158
list(accumulate([10, 20], 100))
159159

160+
def test_accumulate_reenter(self):
161+
# A source iterable (or binop) that calls back into next() on the
162+
# same accumulate object must not be allowed to silently corrupt
163+
# the running total.
164+
class I:
165+
count = 0
166+
def __iter__(self):
167+
return self
168+
def __next__(self):
169+
self.count += 1
170+
if self.count == 2:
171+
return next(it)
172+
return self.count
173+
174+
it = accumulate(I())
175+
self.assertEqual(next(it), 1)
176+
with self.assertRaisesRegex(RuntimeError, "accumulate"):
177+
next(it)
178+
160179
def test_batched(self):
161180
self.assertEqual(list(batched('ABCDEFG', 3)),
162181
[('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)])
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
Fix :func:`itertools.accumulate` silently corrupting its running total
2+
when the source iterable (or the ``func`` callable) re-enters the same
3+
:func:`!accumulate` iterator via a nested call to :func:`!next`. It now
4+
raises :exc:`RuntimeError` on re-entry instead of producing wrong,
5+
silently-dropped values, matching the existing behavior of
6+
:func:`itertools.tee`.

Modules/itertoolsmodule.c

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3079,6 +3079,7 @@ typedef struct {
30793079
PyObject *binop;
30803080
PyObject *initial;
30813081
itertools_state *state;
3082+
int running;
30823083
} accumulateobject;
30833084

30843085
#define accumulateobject_CAST(op) ((accumulateobject *)(op))
@@ -3120,6 +3121,7 @@ itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
31203121
lz->it = it;
31213122
lz->initial = Py_XNewRef(initial);
31223123
lz->state = find_state_by_type(type);
3124+
lz->running = 0;
31233125
return (PyObject *)lz;
31243126
}
31253127

@@ -3160,12 +3162,21 @@ accumulate_next_lock_held(PyObject *op)
31603162
lz->initial = Py_NewRef(Py_None);
31613163
return Py_NewRef(lz->total);
31623164
}
3165+
if (lz->running) {
3166+
PyErr_SetString(PyExc_RuntimeError,
3167+
"cannot re-enter the accumulate iterator");
3168+
return NULL;
3169+
}
3170+
lz->running = 1;
31633171
val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
3164-
if (val == NULL)
3172+
if (val == NULL) {
3173+
lz->running = 0;
31653174
return NULL;
3175+
}
31663176

31673177
if (lz->total == NULL) {
31683178
lz->total = Py_NewRef(val);
3179+
lz->running = 0;
31693180
return lz->total;
31703181
}
31713182

@@ -3174,6 +3185,7 @@ accumulate_next_lock_held(PyObject *op)
31743185
else
31753186
newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
31763187
Py_DECREF(val);
3188+
lz->running = 0;
31773189
if (newtotal == NULL)
31783190
return NULL;
31793191

0 commit comments

Comments
 (0)