Added composed and collect

toolz/functoolz.py

@@ -12,7 +12,7 @@
 
 __all__ = ('identity', 'apply', 'thread_first', 'thread_last', 'memoize',
            'compose', 'compose_left', 'pipe', 'complement', 'juxt', 'do',
-           'curry', 'flip', 'excepts')
+           'curry', 'flip', 'excepts', 'composed')
 
 PYPY = hasattr(sys, 'pypy_version_info')
 
@@ -604,6 +604,49 @@ def compose_left(*funcs):
     return compose(*reversed(funcs))
 
 
+def composed(*func_and_args, **kwargs):
+    """ Decorate a function to create a composition.
+
+    Returns a function that applies the decorated function and then the decorator's argument in a sequence.
+
+    >>> @composed(list)
+    ... def squares(n):
+    ...     for i in range(n):
+    ...         yield i ** 2
+    >>> squares(4) # not a generator anymore
+    [0, 1, 4, 9]
+    >>> @composed(dict)
+    ... def vals_to_squares(values):
+    ...     for v in values:
+    ...         yield v, v ** 2
+    >>> vals_to_squares([3, 2, 8]) # the pairs are gathered in a dict
+    {3: 9, 2: 4, 8: 64}
+    >>> @composed(numpy.stack, axis=-1)  # pass additional arguments
+    ... def animate_noise(n_frames):
+    ...     for i in range(n_frames):
+    ...         noise = numpy.random.uniform if i % 2 else numpy.random.normal
+    ...         yield noise(size=(10, 10))
+    >>> animate_noise(7).shape # stacked along the last axis (axis=-1)
+    (10, 10, 7)
+
+    See Also:
+        compose_left
+    """
+
+    if not func_and_args:
+        raise TypeError('func argument is required')
+    func, *args = func_and_args
+
+    # `compose` can only apply functions with a single argument
+    def final(x):
+        return func(x, *args, **kwargs)
+
+    def decorator(fn):
+        return compose_left(fn, final)
+
+    return decorator
+
+
 def pipe(data, *funcs):
     """ Pipe a value through a sequence of functions
 

toolz/recipes.py

@@ -1,8 +1,8 @@
 import itertools
 from .itertoolz import frequencies, pluck, getter
+from .functoolz import composed
 
-
-__all__ = ('countby', 'partitionby')
+__all__ = ('countby', 'partitionby', 'collect')
 
 
 def countby(key, seq):
@@ -44,3 +44,19 @@ def partitionby(func, seq):
         itertools.groupby
     """
     return map(tuple, pluck(1, itertools.groupby(seq, key=func)))
+
+
+def collect(func):
+    """ Decorate a generator and return a function, that returns a list instead.
+
+    >>> @collect
+    ... def odd_numbers(n):
+    ...     for i in range(n):
+    ...         yield 2 * i + 1
+    >>> odd_numbers(3) # not a generator anymore
+    [1, 3, 5]
+
+    See also:
+        composed
+    """
+    return composed(list)(func)

toolz/tests/test_functoolz.py

@@ -2,7 +2,7 @@
 import toolz
 from toolz.functoolz import (thread_first, thread_last, memoize, curry,
                              compose, compose_left, pipe, complement, do, juxt,
-                             flip, excepts, apply)
+                             flip, excepts, apply, composed)
 from operator import add, mul, itemgetter
 from toolz.utils import raises
 from functools import partial
@@ -679,6 +679,40 @@ def test_compose_left():
         assert compose_left(*compose_left_args)(*args, **kw) == expected
 
 
+def test_composed():
+    for (compose_left_args, args, kw, expected) in generate_compose_left_test_cases():
+        if len(compose_left_args) == 2:
+            func, dec = compose_left_args
+            assert composed(dec)(func)(*args, **kw) == expected
+
+    @composed(dict)
+    def f(n):
+        for i in range(n):
+            yield i, i + 1
+
+    assert f(4) == {0: 1, 1: 2, 2: 3, 3: 4}
+
+    @composed(', '.join)
+    def g(n):
+        for i in range(n):
+            yield str(i)
+
+    assert g(3) == '0, 1, 2'
+
+    # additional args
+    @composed(add, 1)
+    def h():
+        return 2
+
+    assert h() == 3
+
+    @composed(pow, 2)
+    def p():
+        return 5
+
+    assert p() == 25
+
+
 def test_pipe():
     assert pipe(1, inc) == 2
     assert pipe(1, inc, inc) == 3

toolz/tests/test_recipes.py

@@ -1,4 +1,4 @@
-from toolz import first, identity, countby, partitionby
+from toolz import first, identity, countby, partitionby, collect
 
 
 def iseven(x):
@@ -25,3 +25,13 @@ def test_partitionby():
 
     assert ''.join(map(first,
                        partitionby(identity, "Khhhaaaaannnnn!!!!"))) == 'Khan!'
+
+
+def test_collect():
+    @collect
+    def f(n):
+        for i in range(n):
+            if i != 3:
+                yield i ** 2
+
+    assert f(5) == [0, 1, 4, 16]