Skip to content

Commit

Permalink
Make chain a class like the standard sync chain (#103)
Browse files Browse the repository at this point in the history
* Make `chain` a class like the standard sync `chain`

Fixes [GH#102](#102).

Signed-off-by: William So <polyipseity@gmail.com>
  • Loading branch information
polyipseity authored Mar 20, 2023
1 parent cecf2cd commit d7bddaa
Showing 1 changed file with 26 additions and 21 deletions.
47 changes: 26 additions & 21 deletions asyncstdlib/itertools.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,35 +141,40 @@ async def accumulate(iterable, function, *, initial):
yield value


async def chain(*iterables: AnyIterable[T]) -> AsyncIterator[T]:
class chain(AsyncIterator[T]):
"""
An :term:`asynchronous iterator` flattening values from all ``iterables``
The resulting iterator consecutively iterates over and yields all values from
each of the ``iterables``. This is similar to converting all ``iterables`` to
sequences and concatenating them, but lazily exhausts each iterable.
"""
for iterable in iterables:
async with ScopedIter(iterable) as iterator:
async for item in iterator:
yield item


@public_module(__name__, "chain.from_iterable")
async def chain_from_iterable(
iterable: AnyIterable[AnyIterable[T]],
) -> AsyncIterator[T]:
"""
Alternate constructor for :py:func:`~.chain` that lazily exhausts iterables as well
"""
async with ScopedIter(iterable) as iterables:
async for sub_iterable in iterables:
async with ScopedIter(sub_iterable) as iterator:
async for item in iterator:
yield item


chain.from_iterable = chain_from_iterable # type: ignore
__slots__ = ("_impl",)

def __init__(self, *iterables: AnyIterable[T]):
async def impl() -> AsyncIterator[T]:
for iterable in iterables:
async with ScopedIter(iterable) as iterator:
async for item in iterator:
yield item

self._impl = impl()

@staticmethod
async def from_iterable(iterable: AnyIterable[AnyIterable[T]]) -> AsyncIterator[T]:
"""
Alternate constructor for :py:func:`~.chain` that lazily exhausts
iterables as well
"""
async with ScopedIter(iterable) as iterables:
async for sub_iterable in iterables:
async with ScopedIter(sub_iterable) as iterator:
async for item in iterator:
yield item

def __anext__(self) -> Awaitable[T]:
return self._impl.__anext__()


async def compress(
Expand Down

0 comments on commit d7bddaa

Please sign in to comment.