Better way to wait for all tasks from a stream to complete? about aiostream

[vxgmichel]

Hi @Joker-vD, thanks for creating an issue!

I can think of three different ways to implement the feature you're looking for:

import random
import asyncio
from aiostream import stream, pipe


async def do_something(x):
    print(f'Doing something for {x:.2f} s...')
    await asyncio.sleep(x)
    print(f'Done something for {x:.2f} s...')


async def produce_coros():
    for _ in range(5):
        await asyncio.sleep(random.random())
        yield do_something(random.random())


async def _await(arg):
    return await arg


async def main():
    # Version 1
    await stream.combine.amap(produce_coros(), lambda x: x)
    print()
    # Version 2
    await stream.map(produce_coros(), _await)
    print()
    # Version 3
    await stream.flatmap(produce_coros(), stream.just)
    print()


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())
    loop.close()

The last one is probably the simplest, although flatmap is meant to be an advanced operator. Maybe the library is missing an extra helper, but I can't really figure it out. Maybe something like asyncio.as_completed:

from aiostream import stream, pipe, operator

@operator(pipable=True)
def as_completed(source, task_limit=None):
    return stream.flatmap.raw(source, stream.just.raw, task_limit=task_limit)

async def main():
    results = await (as_completed(produce_coros()) | pipe.list())

[vxgmichel]

Now that async_ and await_ have been added to aiostream namespace (see PR #21), the simplest solution is:

from aiostream import stream, await_

async def main():
    await stream.map(produce_coros(), await_)

PR #21 will become available with version 0.3.1.