feat: mapping utils (#78)

a_sync/utils/__init__.py

@@ -1,8 +1,19 @@
 import asyncio
-
+from typing import Awaitable, Dict, Iterator, Mapping, Tuple, TypeVar, List, Callable
+from typing_extensions import ParamSpec
 from a_sync.utils.iterators import (as_yielded, exhaust_iterator,
                                     exhaust_iterators)
 
+T = TypeVar("T")
+KT = TypeVar("KT")
+VT = TypeVar("VT")
+P = ParamSpec("P")
+
+__all__ = [
+    "all",
+    "any",
+    "as_yielded",
+]
 
 async def any(*awaitables) -> bool:
     futs = [asyncio.ensure_future(a) for a in awaitables]
@@ -20,4 +31,18 @@ async def all(*awaitables) -> bool:
             for fut in futs:
                 fut.cancel()
             return False
-    return True
+    return True
+
+async def gather_mapping(mapping: Mapping[KT, Awaitable[VT]]) -> Dict[KT, VT]:
+    results = {k: None for k in mapping.keys()}  # return data in same order
+    async for k, v in as_completed_mapping(mapping):
+        results[k] = v
+    return results
+
+def as_completed_mapping(mapping: Mapping[KT, Awaitable[VT]]) -> Iterator["asyncio.Task[Tuple[KT, VT]]"]:
+    return asyncio.as_completed([__as_completed_wrap(k, v) for k, v in mapping.items()])
+
+async def __as_completed_wrap(k: KT, v: Awaitable[VT]) -> VT:
+    return k, await v
+
+