Fanning out a generator with trio

fanout.py

import trio


async def producer():
    for e in [1, 2, 3]:
        yield e


async def consumer(iterator):
    total = 0
    async for e in iterator:
        print('got {}'.format(e))
        total += e
    return total


async def consumer_wrapper(consumer, its, results, i):
    results[i] = await consumer(its[i])


async def producer_wrapper(producer, its):
    async for e in producer():
        for it in its:
            print('send {}'.format(e))
            await it.send(e)
    for it in its:
        await it.send(None)


class Iterator:
    def __init__(self):
        self.queue = trio.Queue(1)

    async def send(self, e):
        await self.queue.put(e)

    def __aiter__(self):
        return self

    async def __anext__(self):
        next_value = await self.queue.get()
        if next_value is None:
            raise StopAsyncIteration
        return next_value


async def main():
    results = [None, None]
    its = [Iterator(), Iterator()]
    async with trio.open_nursery() as nursery:
        nursery.start_soon(producer_wrapper, producer, its)
        nursery.start_soon(consumer_wrapper, consumer, its, results, 0)
        nursery.start_soon(consumer_wrapper, consumer, its, results, 1)
    print(results)

if __name__ == '__main__':
    trio.run(main)