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)
+flatten = toolz.curry(toolz.flatten)
 get = toolz.curry(toolz.get)
 get_in = toolz.curry(toolz.get_in)
 groupby = toolz.curry(toolz.groupby)

toolz/itertoolz.py

@@ -4,7 +4,7 @@
 import operator
 from functools import partial
 from itertools import filterfalse, zip_longest
-from collections.abc import Sequence
+from collections.abc import Sequence, Mapping
 from toolz.utils import no_default
 
 
@@ -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',
+ 'flatten')
 
 
 def remove(predicate, seq):
@@ -484,6 +485,7 @@ def concat(seqs):
 
  See also:
  itertools.chain.from_iterable equivalent
+ flatten
  """
  return itertools.chain.from_iterable(seqs)
 
@@ -1051,3 +1053,56 @@ def random_sample(prob, seq, random_state=None):
 
  random_state = Random(random_state)
  return filter(lambda _: random_state.random() < prob, seq)
+
+
+def _default_descend(x):
+ return not isinstance(x, (str, bytes, Mapping))
+
+
+def flatten(level, seq, descend=_default_descend):
+ """ Flatten a possibly nested sequence
+
+ Inspired by Javascript's Array.flat(), this is a recursive,
+ depth limited flattening generator. A level 0 flattening will
+ yield the input sequence unchanged. A level -1 flattening
+ will flatten all possible levels of nesting.
+
+ >>> list(flatten(0, [1, [2], [[3]]])) # flatten 0 levels
+ [1, [2], [[3]]]
+ >>> list(flatten(1, [1, [2], [[3]]])) # flatten 1 level
+ [1, 2, [3]]
+ >>> list(flatten(2, [1, [2], [[3]]]))
+ [1, 2, 3]
+ >>> list(flatten(-1, [1, [[[[2]]]]])) # flatten all levels
+ [1, 2]
+
+ An optional ``descend`` function can be provided by the user
+ to determine which iterable objects to recurse into. This function
+ should return a boolean with True meaning it is permissible to descend
+ another level of recursion. The recursion limit of the Python interpreter
+ is the ultimate bounding factor on depth. By default, stings, bytes,
+ and mappings are exempted.
+
+ >>> list(flatten(-1, ['abc', [{'a': 2}, [b'123']]]))
+ ['abc', {'a': 2}, b'123']
+
+ See also:
+ concat
+ """
+ if level < -1:
+ raise ValueError("Level must be >= -1")
+ if not callable(descend):
+ raise ValueError("descend must be a callable boolean function")
+
+ def flat(level, seq):
+ if level == 0:
+ yield from seq
+ return
+
+ for item in seq:
+ if isiterable(item) and descend(item):
+ yield from flat(level - 1, item)
+ else:
+ yield item
+
+ yield from flat(level, seq)

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, flatten)
 from operator import add, mul
 
 
@@ -547,3 +548,37 @@ 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(flatten(0, seq)) == seq
+ assert list(flatten(1, seq)) == seq
+
+ seq = [1, [2, [3]]]
+ assert list(flatten(0, seq)) == seq
+ assert list(flatten(1, seq)) == [1, 2, [3]]
+ assert list(flatten(2, seq)) == [1, 2, 3]
+
+ # Test mappings
+ seq = [{'a': 1}, [1, 2, 3]]
+ assert list(flatten(0, seq)) == seq
+ assert list(flatten(1, seq)) == [{'a': 1}, 1, 2, 3]
+
+ # Test stringsj
+ seq = ["asgf", b"abcd"]
+ assert list(flatten(-1, seq)) == seq
+
+ # Test custom descend function
+ def descend(x):
+ if isinstance(x, str):
+ return len(x) != 1
+ return False
+ seq = ["asdf", [1, 2, 3]]
+ assert list(flatten(1, seq, descend=descend)) == ["a", "s", "d", "f", [1, 2, 3]]
+
+ with pytest.raises(ValueError):
+ list(flatten(0, [1, 2], descend=True))
+
+ with pytest.raises(ValueError):
+ list(flatten(-2, [1, 2]))