feat: map (#113)

BobTheBuidler · Feb 11, 2024 · c44e616 · c44e616
1 parent b1fa768
commit c44e616
Show file tree

Hide file tree

Showing 3 changed files with 23 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,15 @@
+
+from typing import Awaitable, Callable, Literal, Tuple, Union, overload
+
+from a_sync._typing import *
+from a_sync.iter import ASyncIterator
+
+MappingFn = Callable[Concatenate[K, P], Awaitable[V]]
+
+@overload
+def map(coro_fn: MappingFn[K, P, V], *iterables: AnyIterable[K], yields: Literal['keys'], **coro_fn_kwargs: P.kwargs) -> ASyncIterator[K]:...
+@overload
+def map(coro_fn: MappingFn[K, P, V], *iterables: AnyIterable[K], yields: Literal['both'], **coro_fn_kwargs: P.kwargs) -> ASyncIterator[Tuple[K, V]]:...
+def map(coro_fn: MappingFn[K, P, V], *iterables: AnyIterable[K], yields: Literal['keys', 'both'] = 'both', **coro_fn_kwargs: P.kwargs) -> Union[ASyncIterator[K], ASyncIterator[Tuple[K, V]]]:
+    from a_sync.task import TaskMapping
+    return ASyncIterator.wrap(TaskMapping(coro_fn, **coro_fn_kwargs).map(*iterables, yields=yields))