Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 21 additions & 9 deletions Doc/library/itertools.rst
Original file line number Diff line number Diff line change
Expand Up @@ -86,29 +86,38 @@ The following module functions all construct and return iterators. Some provide
streams of infinite length, so they should only be accessed by functions or
loops that truncate the stream.

.. function:: accumulate(iterable[, func])
.. function:: accumulate(iterable[, func, *, initial=None])

Make an iterator that returns accumulated sums, or accumulated
results of other binary functions (specified via the optional
*func* argument). If *func* is supplied, it should be a function
*func* argument).

If *func* is supplied, it should be a function
of two arguments. Elements of the input *iterable* may be any type
that can be accepted as arguments to *func*. (For example, with
the default operation of addition, elements may be any addable
type including :class:`~decimal.Decimal` or
:class:`~fractions.Fraction`.) If the input iterable is empty, the
output iterable will also be empty.
:class:`~fractions.Fraction`.)

Usually, the number of elements output matches the input iterable.
However, if the keyword argument *initial* is provided, the
accumulation leads off with the *initial* value so that the output
has one more element than the input iterable.

Roughly equivalent to::

def accumulate(iterable, func=operator.add):
def accumulate(iterable, func=operator.add, *, initial=None):
'Return running totals'
# accumulate([1,2,3,4,5]) --> 1 3 6 10 15
# accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115
# accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
it = iter(iterable)
try:
total = next(it)
except StopIteration:
return
total = initial
if initial is None:
try:
total = next(it)
except StopIteration:
return
yield total
for element in it:
total = func(total, element)
Expand Down Expand Up @@ -152,6 +161,9 @@ loops that truncate the stream.
.. versionchanged:: 3.3
Added the optional *func* parameter.

.. versionchanged:: 3.8
Added the optional *initial* parameter.

.. function:: chain(*iterables)

Make an iterator that returns elements from the first iterable until it is
Expand Down
6 changes: 6 additions & 0 deletions Lib/test/test_itertools.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,12 @@ def test_accumulate(self):
list(accumulate(s, chr)) # unary-operation
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
self.pickletest(proto, accumulate(range(10))) # test pickling
self.pickletest(proto, accumulate(range(10), initial=7))
self.assertEqual(list(accumulate([10, 5, 1], initial=None)), [10, 15, 16])
self.assertEqual(list(accumulate([10, 5, 1], initial=100)), [100, 110, 115, 116])
self.assertEqual(list(accumulate([], initial=100)), [100])
with self.assertRaises(TypeError):
list(accumulate([10, 20], 100))

def test_chain(self):

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add an optional *initial* argument to itertools.accumulate().
15 changes: 8 additions & 7 deletions Modules/clinic/itertoolsmodule.c.h

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

32 changes: 29 additions & 3 deletions Modules/itertoolsmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -3475,6 +3475,7 @@ typedef struct {
PyObject *total;
PyObject *it;
PyObject *binop;
PyObject *initial;
} accumulateobject;

static PyTypeObject accumulate_type;
Expand All @@ -3484,18 +3485,19 @@ static PyTypeObject accumulate_type;
itertools.accumulate.__new__
iterable: object
func as binop: object = None
*
initial: object = None
Return series of accumulated sums (or other binary function results).
[clinic start generated code]*/

static PyObject *
itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
PyObject *binop)
/*[clinic end generated code: output=514d0fb30ba14d55 input=6d9d16aaa1d3cbfc]*/
PyObject *binop, PyObject *initial)
/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
{
PyObject *it;
accumulateobject *lz;


/* Get iterator. */
it = PyObject_GetIter(iterable);
if (it == NULL)
Expand All @@ -3514,6 +3516,8 @@ itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
}
lz->total = NULL;
lz->it = it;
Py_XINCREF(initial);
lz->initial = initial;
return (PyObject *)lz;
}

Expand All @@ -3524,6 +3528,7 @@ accumulate_dealloc(accumulateobject *lz)
Py_XDECREF(lz->binop);
Py_XDECREF(lz->total);
Py_XDECREF(lz->it);
Py_XDECREF(lz->initial);
Py_TYPE(lz)->tp_free(lz);
}

Expand All @@ -3533,6 +3538,7 @@ accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
Py_VISIT(lz->binop);
Py_VISIT(lz->it);
Py_VISIT(lz->total);
Py_VISIT(lz->initial);
return 0;
}

Expand All @@ -3541,6 +3547,13 @@ accumulate_next(accumulateobject *lz)
{
PyObject *val, *newtotal;

if (lz->initial != Py_None) {
lz->total = lz->initial;
Py_INCREF(Py_None);
lz->initial = Py_None;
Py_INCREF(lz->total);
return lz->total;
}
val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
if (val == NULL)
return NULL;
Expand All @@ -3567,6 +3580,19 @@ accumulate_next(accumulateobject *lz)
static PyObject *
accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
{
if (lz->initial != Py_None) {
PyObject *it;

assert(lz->total == NULL);
if (PyType_Ready(&chain_type) < 0)
return NULL;
it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
lz->initial, lz->it);
if (it == NULL)
return NULL;
return Py_BuildValue("O(NO)O", Py_TYPE(lz),
it, lz->binop?lz->binop:Py_None, Py_None);
}
if (lz->total == Py_None) {
PyObject *it;

Expand Down