Add strong async neutral iteration (#80)

* added any_iter

* added unittests for any_iter

* exporting any_iter at top-level

Author: maxfischer2781

asyncstdlib/__init__.py

@@ -33,7 +33,7 @@
     zip_longest,
     groupby,
 )
-from .asynctools import borrow, scoped_iter, await_each, apply, sync
+from .asynctools import borrow, scoped_iter, await_each, any_iter, apply, sync
 from .heapq import merge, nlargest, nsmallest
 
 __version__ = "3.10.2"
@@ -82,6 +82,7 @@
     "borrow",
     "scoped_iter",
     "await_each",
+    "any_iter",
     "apply",
     "sync",
     # heapq

asyncstdlib/asynctools.py

@@ -380,3 +380,51 @@ async def async_wrapped(*args: Any, **kwargs: Any) -> T:
         return result
 
     return async_wrapped
+
+
+async def any_iter(
+    __iter: Union[
+        Awaitable[AnyIterable[Awaitable[T]]],
+        Awaitable[AnyIterable[T]],
+        AnyIterable[Awaitable[T]],
+        AnyIterable[T],
+    ]
+) -> AsyncIterator[T]:
+    """
+    Provide an async iterator for various forms of "asynchronous iterable"
+
+    Useful to uniformly handle async iterables, awaitable iterables, iterables of
+    awaitables, and similar in an ``async for`` loop. Among other things, this
+    matches all forms of ``async def`` functions providing iterables.
+
+    .. code-block:: python3
+
+        import random
+        import asyncstdlib as a
+
+        # AsyncIterator[T]
+        async def async_iter(n):
+            for i in range(n):
+                yield i
+
+        # Awaitable[Iterator[T]]
+        async def await_iter(n):
+            return [*range(n)]
+
+        some_iter = random.choice([async_iter, await_iter, range])
+        async for item in a.any_iter(some_iter(4)):
+            print(item)
+
+    This function must eagerly resolve each "async layer" before checking if
+    the next layer is as expected. This incurs a performance penalty and
+    non-iterables may be left unusable by this.
+    Prefer :py:func:`~.builtins.iter` to test for iterables with :term:`EAFP`
+    and for performance when only simple iterables need handling.
+    """
+    iterable = __iter if not isinstance(__iter, Awaitable) else await __iter
+    if isinstance(iterable, AsyncIterable):
+        async for item in iterable:
+            yield item if not isinstance(item, Awaitable) else await item
+    else:
+        for item in iterable:
+            yield item if not isinstance(item, Awaitable) else await item

docs/source/api/asynctools.rst

@@ -27,6 +27,11 @@ Async transforming
 
     .. versionadded:: 3.9.3
 
+.. autofunction:: any_iter(iter: (await) (async) iter (await) T)
+    :async-for: :T
+
+    .. versionadded:: 3.10.3
+
 .. autofunction:: await_each(awaitables: iter await T)
     :async-for: :T
 

unittests/test_asynctools.py

@@ -230,7 +230,7 @@ async def check_4(x: int, y: int, z: int) -> int:
 
     t1 = await a.sync(check_3)(x=100)
     t2 = await a.sync(check_4)(x=5, y=5, z=10)
-    t3 = await a.sync(lambda x: x ** 3)(x=5)
+    t3 = await a.sync(lambda x: x**3)(x=5)
 
     with pytest.raises(TypeError):
         a.sync("string")(10)
@@ -252,3 +252,36 @@ async def coro():
         return coro()
 
     assert await nocoro_async(5) == 5
+
+
+async def await_iter(n: int):
+    return [*range(n)]
+
+
+async def async_iter(n: int):
+    for i in range(n):
+        yield i
+
+
+async def await_value(i):
+    return i
+
+
+async def await_iter_await(n: int):
+    return [await_value(i) for i in range(n)]
+
+
+async def await_async_iter_await(n: int):
+    for i in range(n):
+        yield await_value(i)
+
+
+@pytest.mark.parametrize("n", [0, 1, 12])
+@pytest.mark.parametrize(
+    "any_iterable_t",
+    [range, await_iter, async_iter, await_iter_await, await_async_iter_await],
+)
+@sync
+async def test_any_iter(n, any_iterable_t):
+    iterable = any_iterable_t(n)
+    assert [item async for item in a.any_iter(iterable)] == [*range(n)]