What is the most convenient way to map a Coroutine for asynchronous execution in Python?

Viewed 681

I am trying to "map" an Iterable asynchronously over multiple functions in a Python project. I have found a few ways to do this, namely:

  • Method 1: asyncio.gather (this works).
  • Method 2: A regular for loop (this does not work, shown for reference).
  • Method 3: A decorator that accepts an Iterable as input along with the name of an argument for injection (this works too and is equivalent to Method 1).

Below is the code:

from asyncio import AbstractEventLoop
from datetime import datetime
from functools import wraps
from typing import AsyncIterable, Coroutine, Iterable, List, Union
import asyncio
import inspect
import os


ITERATIONS: List[int] = list(range(3))


def map_async(
        iterable: Union[AsyncIterable, Iterable],
        kwarg: str,
    ) -> Coroutine:
    """
    A decorator that maps an iterable for asynchronous
    consumption by a Coroutine.

    Parameters
    ----------
    iterable : Union[AsyncIterable, Iterable]
        An iterable (either sync or async).
    kwarg : str
        Name of the argument to inject into the wrapped Coroutine
        for each iteration of the iterable.
    """
    def wrapper(async_func: Coroutine):
        @wraps(async_func)
        async def wrapped(*args, **kwargs):
            if inspect.isasyncgen(iterable):
                return await asyncio.gather(
                    *[
                        async_func(*args, **{kwarg: xx}, **kwargs)
                        async for xx in iterable
                    ]
                )
            else:
                return await asyncio.gather(
                    *(
                        async_func(*args, **{kwarg: xx}, **kwargs)
                        for xx in iterable
                    )
                )
        return wrapped
    return wrapper


# Method 1 (a):
async def foo(count: int, sleep_duration: int):
    await asyncio.sleep(sleep_duration)
    print(count)


# Method 1 (b):
async def async_foo_gather(sleep_duration: int):
    await asyncio.gather(*(foo(count, sleep_duration) for count in ITERATIONS))


# Method 2:
async def async_foo_for(sleep_duration: int):
    for count in ITERATIONS:
        await foo(count, sleep_duration)


# Method 3:
@map_async(iterable=ITERATIONS, kwarg='count')
async def async_foo_mapped(sleep_duration: int, **kwargs):
    await asyncio.sleep(sleep_duration)
    print(kwargs['count'])


if __name__ == '__main__':

    loop: AbstractEventLoop = asyncio.get_event_loop()

    TERMINAL_SIZE = os.get_terminal_size()

    hrule: str = '-' * TERMINAL_SIZE.columns

    print(hrule)

    # Method 1: asyncio.gather over an Iterable.

    print('async_foo_gather')
    start: datetime = datetime.now()
    loop.run_until_complete(async_foo_gather(sleep_duration=5))
    stop: datetime = datetime.now()
    print('DURATION: ', (stop - start).total_seconds())

    print(hrule)

    # Method 2: synchronous for loop.

    print('async_foo_for')
    start: datetime = datetime.now()
    loop.run_until_complete(async_foo_for(sleep_duration=5))
    stop: datetime = datetime.now()
    print('DURATION: ', (stop - start).total_seconds())

    print(hrule)

    # Method 3: map_async decorator.

    print('async_foo_mapped')
    start: datetime = datetime.now()
    loop.run_until_complete(async_foo_mapped(sleep_duration=5))
    stop: datetime = datetime.now()
    print('DURATION: ', (stop - start).total_seconds())

    print(hrule)

This results in the following output:

---------------------------------------------------------------------------------------------------------------------------------------
async_foo_gather
0
1
2
DURATION:  5.000184
---------------------------------------------------------------------------------------------------------------------------------------
async_foo_for
0
1
2
DURATION:  15.00077
---------------------------------------------------------------------------------------------------------------------------------------
async_foo_mapped
0
1
2
DURATION:  5.001652
---------------------------------------------------------------------------------------------------------------------------------------

Method 1 and Method 3 are equivalent but I feel like there must be a more elegant or idiomatic way to do this?

Helpers such as the map_async decorator etc could live in their own module but, overall, I'd really appreciate something that saves having to define two functions in each usage (as in Method 1, nested or otherwise) because this transformation is needed in many places.

Thanks

0 Answers
Related