rainbowadn/rainbowadn/core/gather.py
2022-06-23 00:51:43 +03:00

36 lines
622 B
Python

import asyncio
from typing import AsyncIterable, TypeVar
__all__ = ('gather', 'asum', 'alist', 'set_gather_asyncio', 'set_gather_linear',)
_gather = asyncio.gather
async def _local_gather(*args):
return [await arg for arg in args]
def set_gather_asyncio():
global _gather
_gather = asyncio.gather
def set_gather_linear():
global _gather
_gather = _local_gather
def gather(*args):
return _gather(*args)
async def asum(iterable):
return sum(await gather(*iterable))
T = TypeVar('T')
async def alist(aiterable: AsyncIterable[T]) -> list[T]:
return [x async for x in aiterable]