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
forloop (this does not work, shown for reference). - Method 3: A decorator that accepts an
Iterableas 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