evhub / coconut

Simple, elegant, Pythonic functional programming.
http://coconut-lang.org
Apache License 2.0
4.04k stars 120 forks source link

Add an `async_map` function utilizing `anyio` #795

Closed evhub closed 9 months ago

evhub commented 9 months ago

Currently, as of #124, fmap can map over asynchronous iterables, but we don't currently have a built-in way to map an async function over a normal iterable.

Implementation:

async def async_map[T, U](
    async_func: async T -> U,
    *iters: T$[],
    strict: bool = False
) -> U[]:
    """Map async_func over iters asynchronously using anyio."""
    import anyio
    results = []
    async def store_func_in_of(i, args):
        got = await async_func(*args)
        results.extend([None] * (1 + i - len(results)))
        results[i] = got
    async with anyio.create_task_group() as nursery:
        for i, args in enumerate(zip(*iters, strict=strict)):
            nursery.start_soon(store_func_in_of, i, args)
    return results