diff --git a/a_sync/__init__.py b/a_sync/__init__.py index c67cec6b..f914bfcc 100644 --- a/a_sync/__init__.py +++ b/a_sync/__init__.py @@ -8,9 +8,7 @@ from a_sync.primitives import * from a_sync.singleton import ASyncGenericSingleton from a_sync.task import TaskMapping, create_task -from a_sync.utils import all, any, as_yielded -from a_sync.utils.as_completed import as_completed -from a_sync.utils.gather import gather +from a_sync.utils import all, any, as_completed, as_yielded, gather, map # I alias the aliases for your convenience. # I prefer "aka" but its meaning is not intuitive when reading code so I created both aliases for you to choose from. @@ -30,6 +28,7 @@ "exhaust_iterator", "exhaust_iterators", "gather", + "map", "ASyncIterable", "ASyncIterator", "ASyncGenericSingleton", diff --git a/a_sync/utils/__init__.py b/a_sync/utils/__init__.py index 3010c0f2..c8faa34f 100644 --- a/a_sync/utils/__init__.py +++ b/a_sync/utils/__init__.py @@ -2,12 +2,18 @@ from a_sync.utils.iterators import (as_yielded, exhaust_iterator, exhaust_iterators) +from a_sync.utils.as_completed import as_completed +from a_sync.utils.gather import gather +from a_sync.utils.map import map __all__ = [ "all", "any", + "as_completed", "as_yielded", + "gather", + "map", "exhaust_iterator", "exhaust_iterators", ] diff --git a/a_sync/utils/map.py b/a_sync/utils/map.py new file mode 100644 index 00000000..60efea97 --- /dev/null +++ b/a_sync/utils/map.py @@ -0,0 +1,13 @@ + +from typing import Awaitable, Callable, Literal, Tuple, Union, overload + +from a_sync._typing import K, P, V, AnyIterable +from a_sync.iter import ASyncIterator +from a_sync.task import TaskMapping + +@overload +def map(coro_fn: Callable[P, Awaitable[V]], *iterables: AnyIterable[P.args], yields: Literal['keys'], **coro_fn_kwargs: P.kwargs) -> ASyncIterator[K]:... +@overload +def map(coro_fn: Callable[P, Awaitable[V]], *iterables: AnyIterable[P.args], yields: Literal['both'], **coro_fn_kwargs: P.kwargs) -> ASyncIterator[Tuple[K, V]]:... +def map(coro_fn: Callable[P, Awaitable[V]], *iterables: AnyIterable[P.args], yields: Literal['keys', 'both'] = 'both', **coro_fn_kwargs: P.kwargs) -> Union[ASyncIterator[K], ASyncIterator[Tuple[K, V]]]: + return ASyncIterator.wrap(TaskMapping(coro_fn, **coro_fn_kwargs).map(*iterables, yields=yields)) \ No newline at end of file