イテレータを分割する

itertoolsには、Trueになる要素を返すifilterと、Falseになる要素を返すifilterfalseがありますが、ifilterとifilterfalseの両方を返す関数があったら便利じゃないか？
（schemeやclojureなどにはあるようです）

あるいは、True/Falseだけでなく、任意の値でイテレータを分割できたら便利ではないか？と考え、作って見ました。

いえ、実際に使いでがあるかどうかは知らないのですが、一応。

__metaclass__ = type
from itertools import (
    ifilter,
    ifilterfalse,
    tee,
)

def separate(iterable, pred):
    """
    >>> it1, it2 = separate(xrange(10), lambda x: x % 3 == 0)
    >>> list(it1)
    [0, 3, 6, 9]
    >>> list(it2)
    [1, 2, 4, 5, 7, 8]
    """

    it1, it2 = tee(iterable)
    return (ifilter(pred, it1), ifilterfalse(pred, it2))


def separate_n(iterable, n, key):
    """
    >>> it0, it1, it2 = separate_n(xrange(10), 3, lambda x:x % 3)
    >>> list(it0)
    [0, 3, 6, 9]
    >>> list(it1)
    [1, 4, 7]
    >>> list(it2)
    [2, 5, 8]
    """

    return tuple(ifilter(lambda x, i=i:key(x) == i, it)
                      for i, it in enumerate(tee(iterable, n)))


class Tee:
    def __init__(self, it, L=None):
        self._i = -1
        self._it = iter(it)
        if L is not None:
            self._L = L
        else:
            self._L = []
    
    def __iter__(self):
        return self
    
    def next(self):
        self._i += 1
        if 0 <= self._i < len(self._L):
            return self._L[self._i]
        else:
            val = self._it.next()
            self._L.append(val)
            return val

    def tee(self):
        return Tee(self._it, self._L)


class Separator:
    """
    >>> s = Separator(xrange(10), lambda x:x%3)
    >>> list(s.filter(0))
    [0, 3, 6, 9]
    >>> list(s.filter(1))
    [1, 4, 7]
    >>> list(s.filter(2))
    [2, 5, 8]
    >>> list(s.filter(3))
    []
    >>> list(s.filter(0))
    [0, 3, 6, 9]
    """
    def __init__(self, iterable, key):
        self._key = key
        self._iterable = Tee(iterable)        

    def filter(self, key):
        return ifilter(lambda x:self._key(x) == key, self._iterable.tee())

def main():
    import doctest
    doctest.testmod()

if __name__ == "__main__":
    main()

ただし、ジェネレータの利点は、巨大なリストを生成しないで済むことですが、
Teeやitertools.teeはが既にiterableの以前の値を保存するので、itertools.countのような無限ジェネレータに使うと、メモリの無駄になる可能性があります。