Sugar above the itertools module.

#!/usr/bin/env python
# -*-coding: utf8-*-
# Author: Gribouillis for the python forum at www.daniweb.com
# Date: October 2012
# License: Public Domain
# Use this code freely
from __future__ import unicode_literals, print_function

#== docstring ==========================================
__doc__ = """ module: itergadgets -

    This module adds a thin layer of sugar on top of
    the itertools module's api:
        
        * Many functions in itertools take a function
        and a sequence as an argument. This module
        transforms these functions into decorators,
        for example if 'func' is a function which applies
        to a single item of an iterable, then all of
        
        mapper(func)
        truefilter(func)
        falsefilter(func)
        whiledropper(func)
        whiletaker(func)
        starmapper(func)
        grouper(func)
        sorter_grouper(func)
        
        are callables which take a single iterable as
        argument and return an iterable.
    
        * An unpleasant fact about itertools.groupby is
        that it yields pairs (key, group) instead of just
        groups. This module's grouper(function) yields
        a sequence of group object with a .key attribute
        instead
        
        * This module also contains a few utility functions
        
        ppartial()
        automaton()
        compose()
        consume()
        identity()
        powers()
        returner()
        slicer()
        ungrouped()
        
        see each function's documentation.
"""
#== imports ============================================

from functools import partial
import itertools

#== body ===============================================

def _make_one(*args, **kwd):
    def decorator(cls):
        name = cls.__name__
        globals()["_" + name] = cls
        return cls(*args, **kwd)
    return decorator


def ppartial(*args, **kwd):
    return partial(partial, *args, **kwd)

class _PPartial(partial):
    def __new__(cls, func, **kwd):
        instance = partial.__new__(cls, partial, func)
        instance.__dict__.update(kwd)
        return instance

class automaton(object):
    """automaton(binary function) --> automaton object
    
    Return a callable such that if F = automaton(f), then
        F(val, (x, y, z, ...)) is the sequence
        (val, f(val, x), f(f(val, x), y), ...)
        
    Similar to Haskell's scanl.
    
    Example
    
        >>> @automaton
        ... def states(state, letter):
        ...   return state + (1+len(state)) * letter
        >>> list( states("", "abcd") )
        [u'', u'a', u'abb', u'abbcccc', u'abbccccdddddddd']
    """
    def __init__(self, func):
        self.func = func
    def __call__(self, state, iterable):
        func = self.func
        yield state
        for transition in iterable:
            state = func(state, transition)
            yield state

class compose(object):
    """compose(f0, f1, f2, ...) --> callable
    
    return a callable F such that F(*args, **kwd) is equivalent to
        ...f2(f1(f0(*args, **kwd)))
        
    Example
        
        >>> f = compose(lambda x: x**2, lambda x: x + 1)
        >>> [f(x) for x in range(5)]
        [1, 2, 5, 10, 17]
    """
    def __init__(self, *functions):
        self.functions = functions if functions else (
                                returner(itertools.repeat(None)),)
    def __call__(self, *args, **kwd):
        value = self.functions[0](*args, **kwd)
        for f in self.functions[1:]:
            value = f(value)
        return value

def consume(sequence):
    """consume an iterable, returning the last value.
    
    If the iterable is empty, None is returned.
    
    Example
    
        >>> consume(xrange(5))
        4
        
    """
    x = None
    for x in sequence:
        pass
    return x

@_make_one(itertools.ifilter)
class truefilter(_PPartial):
    """truefilter(function or None) --> partial object
    
    Return an callable F such that
        F(iterable) <=> itertools.ifilter(function, iterable)
        
    example
    
        >>> @truefilter
        ... def odd(item):
        ...   return item % 2
        >>>
        >>> list( odd(range(10)) )
        [1, 3, 5, 7, 9]
        
    """

@_make_one(itertools.ifilterfalse)
class falsefilter(_PPartial):
    """falsefilter(function or None) --> partial object
    
    Return an callable F such that
        F(iterable) <=> itertools.ifilterfalse(function, iterable)
        
    example
    
        >>> @falsefilter
        ... def drop_odd(item):
        ...   return item % 2
        >>>
        >>> list( drop_odd(range(10)) )
        [0, 2, 4, 6, 8]
        
    """

@_make_one(itertools.dropwhile)
class whiledropper(_PPartial):
    """whiledropper(predicate) --> partial object
    
    Return an callable F such that
        F(iterable) <=> itertools.dropwhile(predicate, iterable)
        
    example
    
        >>> @whiledropper
        ... def drop_head(item):
        ...   return item < 5
        >>> list( drop_head(range(10)) )
        [5, 6, 7, 8, 9]
        
    """

def grouper(func):
    """grouper(function) --> callable
    
    Return a callable F such that F(iterable) returns a sequence
    of sub-iterables (groups). This is similar to
    itertools.groupby(iterable, key = function) except that a sequence
    of groups is generated instead of a sequence of pairs (key, group),
    and each group has a .key attribute containing the key's value.
    
    Example
        
        >>> @grouper
        ... def quot3(item):
        ...   return item // 3
        >>> list( list(g) for g in quot3(xrange(10)) )
        [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
        >>> list( g.key for g in quot3(xrange(10)) )
        [0, 1, 2, 3]
    
    """
    return compose(partial(itertools.groupby, key = func), mapper(Group))

def identity(x):
    """Identity function, returns its argument.
    
        >>> identity(5)
        5
    """
    return x

@_make_one(itertools.imap)
class mapper(_PPartial):
    """mapper(function) --> partial object
    
    Return an callable F such that
        F(iterable) <=> itertools.imap(function, iterable)
        
    example
    
        >>> @mapper
        ... def double(item):
        ...   return item * 2
        >>>
        >>> list( double(range(5)) )
        [0, 2, 4, 6, 8]
        
    """

class powers(object):
    """powers(function) --> powers object
    
    Return a callable such that
        powers(f)(x) is the sequence (x, f(x), f(f(x)), f(f(f(x))), ...) 
    
    example
    
        >>> @powers
        ... def hailstones(x):
        ...   if x == 1:
        ...     raise StopIteration
        ...   return ((3*x + 1) if x % 2 else x) // 2
        >>>
        >>> list( hailstones(13) )
        [13, 20, 10, 5, 8, 4, 2, 1]
        
    """
    def __init__(self, func):
        self.func = func
    def __call__(self, value):
        func = self.func
        while True:
            yield value
            value = func(value)

class returner(object):
    """returner(iterable) --> callable
    
    return a function F such that if F = returner((v0, v1, v2, ...))
    then the first call to F(*args, **kwd) will return v0, the second
    call will return v1, etc (independently of the arguments).

    Example
    
        >>> f = returner( xrange(10, 13) )
        >>> [f("spam") for i in xrange(3) ]
        [10, 11, 12]
    """
    def __init__(self, iterable):
        self.seq = iter(iterable)
    def __call__(self, *args, **kwd):
        return next(self.seq)

class slicer(object):
    """slicer([start], stop, [step]) --> slicer object
    
    Return a callable F such that
        F(iterable) <=> itertools.islice(iterable , start, stop, step)
    
    example
    
        >>> list( slicer(1, 5)(range(10, 20)) )
        [11, 12, 13, 14]
        
    """
    def __init__(self, *args):
        self.args = args
    def __call__(self, seq):
        return itertools.islice(seq, *self.args)

shift = slicer(1, None)

@_make_one(itertools.starmap)
class starmapper(_PPartial):
    """starmapper(function) --> partial object
    
    Return an callable F such that
        F(iterable) <=> itertools.starmap(function, iterable)
        
    example
    
        >>> @starmapper
        ... def square(x, y):
        ...   return x ** 2 + y ** 2
        >>> list( square(zip(range(3), range(10,13)) ))
        [100, 122, 148]
    
    """
    

def sorter_grouper(function):
    """sorter_grouper(function) --> callable
    
    similar to grouper(function), but the returned callable
    sorts its iterable argument on the keys before grouping.
    
    Example
        
        >>> @sorter_grouper
        ... def mod3(item):
        ...   return item % 3
        >>> list( list(g) for g in mod3(xrange(10)) )
        [[0, 3, 6, 9], [1, 4, 7], [2, 5, 8]]
        >>> list( g.key for g in mod3(xrange(10)) )
        [0, 1, 2]
    """
    return compose(partial(sorted, key = function), grouper(function))

@_make_one(itertools.takewhile)
class whiletaker(_PPartial):
    """whiletaker(predicate) --> partial object
    
    Return an callable F such that
        F(iterable) <=> itertools.takewhile(predicate, iterable)
        
    example
    
        >>> @whiletaker
        ... def take_some(item):
        ...   return item < 5
        >>>
        >>> list( take_some(range(10)) )
        [0, 1, 2, 3, 4]
        
    """

class Group(object):
    """A wrapper around itertools._grouper instances
    
    This is the type of the groups generated by grouper().
    Each group is an iterable, with an attribute group.key.
    These objects have the limitations described
    in itertools.groupby()'s documentation.
    """
    __slots__ = ("_grouper", "key")
    
    def __init__(self, pair):
        self.key, self._grouper = pair
        
    def __iter__(self):
        return self._grouper

ungrouped = itertools.chain.from_iterable

#== script code  =======================================

if __name__ == "__main__":
    import doctest
    doctest.testmod()

#=======================================================