Add itertoolz.flatten

toolz/curried/__init__.py

@@ -65,6 +65,7 @@
 drop = toolz.curry(toolz.drop)
 excepts = toolz.curry(toolz.excepts)
 filter = toolz.curry(toolz.filter)
+flat = toolz.curry(toolz.flat)
 get = toolz.curry(toolz.get)
 get_in = toolz.curry(toolz.get_in)
 groupby = toolz.curry(toolz.groupby)

toolz/itertoolz.py

@@ -13,7 +13,8 @@
            'first', 'second', 'nth', 'last', 'get', 'concat', 'concatv',
            'mapcat', 'cons', 'interpose', 'frequencies', 'reduceby', 'iterate',
            'sliding_window', 'partition', 'partition_all', 'count', 'pluck',
-           'join', 'tail', 'diff', 'topk', 'peek', 'peekn', 'random_sample')
+           'join', 'tail', 'diff', 'topk', 'peek', 'peekn', 'random_sample',
+           'flat')
 
 
 def remove(predicate, seq):
@@ -1055,3 +1056,14 @@ def random_sample(prob, seq, random_state=None):
 
         random_state = Random(random_state)
     return filter(lambda _: random_state.random() < prob, seq)
+
+
+def flat(level, seq):
+    """ Flatten a possible nested sequence by n levels """
+    if level < 0:
+        raise ValueError("level must be >= 0")
+    for item in seq:
+        if level == 0 or not hasattr(item, '__iter__'):
+            yield item
+        else:
+            yield from flat(level - 1, item)

toolz/tests/test_itertoolz.py

@@ -1,3 +1,4 @@
+import pytest
 import itertools
 from itertools import starmap
 from toolz.utils import raises
@@ -13,7 +14,7 @@
                              reduceby, iterate, accumulate,
                              sliding_window, count, partition,
                              partition_all, take_nth, pluck, join,
-                             diff, topk, peek, peekn, random_sample)
+                             diff, topk, peek, peekn, random_sample, flat)
 from operator import add, mul
 
 
@@ -547,3 +548,17 @@ def test_random_sample():
     assert mk_rsample(b"a") == mk_rsample(u"a")
 
     assert raises(TypeError, lambda: mk_rsample([]))
+
+
+def test_flat():
+    seq = [1, 2, 3, 4]
+    assert list(flat(0, seq)) == seq
+    assert list(flat(1, seq)) == seq
+
+    seq = [1, [2, [3]]]
+    assert list(flat(0, seq)) == seq
+    assert list(flat(1, seq)) == [1, 2, [3]]
+    assert list(flat(2, seq)) == [1, 2, 3]
+
+    with pytest.raises(ValueError):
+        list(flat(-1, seq))