lazy_categorize.py

import collections


def lazy_categorize(iterable, predicate):
    a, b = collections.deque(), collections.deque()
    empty = False

    def _gen_generator(t=True):
        d = a if t else b
        d_inv = b if t else a

        def g():
            nonlocal empty
            while not empty:
                if len(d) > 0:
                    yield d.popleft()
                else:
                    try:
                        e = next(iterable)
                        p = predicate(e)
                        if not t ^ p:
                            yield e
                        else:
                            d_inv.append(e)
                    except StopIteration:
                        empty = True

        return g()

    return (_gen_generator(True), _gen_generator(False))


if __name__ == '__main__':
    r = iter(range(10000))
    t, f = lazy_categorize(r, lambda x: x % 2 == 0)

    for _ in range(25):
        print("evens ", next(t))
    for _ in range(100):
        print("odds ", next(f))
    for _ in range(100):
        print("range at ", next(r))
    for _ in range(25):
        print("evens ", next(t))