81 lines
1.6 KiB
Python
81 lines
1.6 KiB
Python
import asyncio
|
|
from typing import Any, AsyncIterable, Awaitable, Coroutine, TypeVar, overload
|
|
|
|
__all__ = ('gather', 'asum', 'alist', 'set_gather_asyncio', 'set_gather_linear', 'aidentity',)
|
|
|
|
_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
|
|
|
|
|
|
T0 = TypeVar('T0')
|
|
T1 = TypeVar('T1')
|
|
T2 = TypeVar('T2')
|
|
T3 = TypeVar('T3')
|
|
T4 = TypeVar('T4')
|
|
|
|
|
|
@overload
|
|
def gather(
|
|
a0: Coroutine[Any, Any, T0],
|
|
a1: Coroutine[Any, Any, T1],
|
|
) -> Awaitable[tuple[T0, T1]]: ...
|
|
|
|
|
|
@overload
|
|
def gather(
|
|
a0: Coroutine[Any, Any, T0],
|
|
a1: Coroutine[Any, Any, T1],
|
|
a2: Coroutine[Any, Any, T2],
|
|
) -> Awaitable[tuple[T0, T1, T2]]: ...
|
|
|
|
|
|
@overload
|
|
def gather(
|
|
a0: Coroutine[Any, Any, T0],
|
|
a1: Coroutine[Any, Any, T1],
|
|
a2: Coroutine[Any, Any, T2],
|
|
a3: Coroutine[Any, Any, T3],
|
|
) -> Awaitable[tuple[T0, T1, T2, T3]]: ...
|
|
|
|
|
|
@overload
|
|
def gather(
|
|
a0: Coroutine[Any, Any, T0],
|
|
a1: Coroutine[Any, Any, T1],
|
|
a2: Coroutine[Any, Any, T2],
|
|
a3: Coroutine[Any, Any, T3],
|
|
a4: Coroutine[Any, Any, T4],
|
|
) -> Awaitable[tuple[T0, T1, T2, T3, T4]]: ...
|
|
|
|
|
|
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]
|
|
|
|
|
|
async def aidentity(value: T) -> Coroutine[Any, Any, T]:
|
|
return value
|