itertools: Add accumulate function

itertools/itertools.py

@@ -55,3 +55,14 @@ def tee(iterable, n=2):
 def starmap(function, iterable):
     for args in iterable:
         yield function(*args)
+
+def accumulate(iterable, func=lambda x, y: x + y):
+    it = iter(iterable)
+    try:
+        acc = next(it)
+    except StopIteration:
+        return
+    yield acc
+    for element in it:
+        acc = func(acc, element)
+        yield acc

itertools/test_itertools.py

@@ -14,3 +14,9 @@ assert list(itertools.islice(itertools.cycle([1, 2, 3]), 10)) == [1, 2, 3, 1, 2,
 assert list(itertools.islice(itertools.cycle(reversed([1, 2, 3])), 7)) == [3, 2, 1, 3, 2, 1, 3]
 
 assert list(itertools.starmap(lambda x, y: x * y, [[1, 2], [2, 3], [3, 4]])) == [2, 6, 12]
+
+assert list(itertools.accumulate([])) == []
+assert list(itertools.accumulate([0])) == [0]
+assert list(itertools.accumulate([0, 2, 3])) == [0, 2, 5]
+assert list(itertools.accumulate(reversed([0, 2, 3]))) == [3, 5, 5]
+assert list(itertools.accumulate([1, 2, 3], lambda x, y: x * y)) == [1, 2, 6]