feat: map util

BobTheBuidler · Feb 11, 2024 · c9ac174 · c9ac174
1 parent 3e4a95e
commit c9ac174
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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))