How to split one iterator into two?

Question

How would you split one iterator into two without iterating twice or using additional memory to store all the data?

Solution when you can store everything in memory:

l = [{'a': i, 'b': i * 2} for i in range(10)]
def a(iterator):
    for item in iterator:
        print(item)
def b(iterator):
   for item in iterator:
        print(item)

a([li['a'] for li in l])
b([li['b'] for li in l])

or if you can iterate twice,

class SomeIterable(object):
    def __iter__(self):
        for i in range(10):
            yield {'a': i, 'b': i * 2}


def a(some_iterator):
    for item in some_iterator:
        print(item)


def b(some_iterator):
    for item in some_iterator:
        print(item)


s = SomeIterable()

a((si['a'] for si in s))
b((si['b'] for si in s))

But how would I make it if I just want to iterate once?

Answer 1

From the clarification in the comments, a and b are external library functions you cannot rewrite, but it's okay to interleave their execution. In that case, what you want is possible, but it pretty much requires threads:

import multiprocessing.pool # for ThreadPool, not multiprocessing
import Queue

_endofinput = object()

def _queueiter(queue):
    while True:
        item = queue.get()
        if item is _endofinput:
            break
        yield item

def parallel_execute(funcs, iterable, maxqueue):
    '''Interleaves the execution of funcs[0](iterable), funcs[1](iterable), etc.

    No function is allowed to lag more than maxqueue items behind another.
    (This will require adjustment if a function might return before consuming
    all input.)

    Makes only one pass over iterable.

    '''

    queues = [Queue.Queue(maxsize=maxqueue) for func in funcs]
    queueiters = [_queueiter(queue) for queue in queues]
    threadpool = multiprocessing.pool.ThreadPool(processes=len(funcs))

    results = threadpool.map_async(lambda (f, x): f(x), zip(funcs, queueiters))

    for item in iterable:
        for queue in queues:
            queue.put(item)

    for queue in queues:
        queue.put(_endofinput)

    threadpool.close()
    return results.get()

Answer 2

If the functions that consume your two iterators are not under your control and don't return control of the program to your code before consuming all of the iterator contents, there is no way to do what you want. You'll either need to hold all of the data in memory in between function calls or regenerate the iterator for the second function.

Now, if your functions were generators (that yield back to your code after consuming some small number of items from the input), you could make it work with itertools.tee. There might also be some other partial workarounds if you can call one or both of your functions with various parts of the input data at a time and then somehow compile the results of the repeated calls together into the desired output. Otherwise you're probably out of luck

Answer 3

Ok, if your functions are stateless, but still expect an iterable as argument, and that's the whole problem, then this should do:

for si in s:
    a([si['a']])
    b([si['b']])