chore: 添加虚拟环境到仓库
- 添加 backend_service/venv 虚拟环境 - 包含所有Python依赖包 - 注意:虚拟环境约393MB,包含12655个文件
This commit is contained in:
@@ -0,0 +1,46 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
"""
|
||||
itertools and builtins for AsyncIO and mixed iterables
|
||||
"""
|
||||
|
||||
__author__ = "Amethyst Reese"
|
||||
from . import asyncio
|
||||
from .__version__ import __version__
|
||||
from .builtins import (
|
||||
all,
|
||||
any,
|
||||
enumerate,
|
||||
iter,
|
||||
list,
|
||||
map,
|
||||
max,
|
||||
min,
|
||||
next,
|
||||
set,
|
||||
sum,
|
||||
tuple,
|
||||
zip,
|
||||
)
|
||||
from .itertools import (
|
||||
accumulate,
|
||||
batched,
|
||||
chain,
|
||||
combinations,
|
||||
combinations_with_replacement,
|
||||
compress,
|
||||
count,
|
||||
cycle,
|
||||
dropwhile,
|
||||
filterfalse,
|
||||
groupby,
|
||||
islice,
|
||||
permutations,
|
||||
product,
|
||||
repeat,
|
||||
starmap,
|
||||
takewhile,
|
||||
tee,
|
||||
zip_longest,
|
||||
)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
This file is automatically generated by attribution.
|
||||
|
||||
Do not edit manually. Get more info at https://attribution.omnilib.dev
|
||||
"""
|
||||
|
||||
__version__ = "0.13.0"
|
||||
@@ -0,0 +1,253 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
"""
|
||||
Friendlier version of asyncio standard library.
|
||||
|
||||
Provisional library. Must be imported as `aioitertools.asyncio`.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Iterable
|
||||
from typing import Any, cast, Optional
|
||||
|
||||
from .builtins import iter as aiter, maybe_await
|
||||
from .types import AnyIterable, AsyncIterator, MaybeAwaitable, T
|
||||
|
||||
|
||||
async def as_completed(
|
||||
aws: Iterable[Awaitable[T]],
|
||||
*,
|
||||
timeout: Optional[float] = None,
|
||||
) -> AsyncIterator[T]:
|
||||
"""
|
||||
Run awaitables in `aws` concurrently, and yield results as they complete.
|
||||
|
||||
Unlike `asyncio.as_completed`, this yields actual results, and does not require
|
||||
awaiting each item in the iterable.
|
||||
|
||||
Cancels all remaining awaitables if a timeout is given and the timeout threshold
|
||||
is reached.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in as_completed(futures):
|
||||
... # use value immediately
|
||||
|
||||
"""
|
||||
done: set[Awaitable[T]] = set()
|
||||
pending: set[Awaitable[T]] = {asyncio.ensure_future(a) for a in aws}
|
||||
remaining: Optional[float] = None
|
||||
|
||||
if timeout and timeout > 0:
|
||||
threshold = time.time() + timeout
|
||||
else:
|
||||
timeout = None
|
||||
|
||||
while pending:
|
||||
if timeout:
|
||||
remaining = threshold - time.time()
|
||||
if remaining <= 0:
|
||||
for fut in pending:
|
||||
if isinstance(fut, asyncio.Future):
|
||||
fut.cancel()
|
||||
else: # pragma: no cover
|
||||
pass
|
||||
raise asyncio.TimeoutError()
|
||||
|
||||
# asyncio.Future inherits from typing.Awaitable
|
||||
# asyncio.wait takes Iterable[Union[Future, Generator, Awaitable]], but
|
||||
# returns Tuple[Set[Future], Set[Future]. Because mypy doesn't like assigning
|
||||
# these values to existing Set[Awaitable] or even Set[Union[Awaitable, Future]],
|
||||
# we need to first cast the results to something that we can actually use
|
||||
# asyncio.Future: https://github.com/python/typeshed/blob/72ff7b94e534c610ddf8939bacbc55343e9465d2/stdlib/3/asyncio/futures.pyi#L30
|
||||
# asyncio.wait(): https://github.com/python/typeshed/blob/72ff7b94e534c610ddf8939bacbc55343e9465d2/stdlib/3/asyncio/tasks.pyi#L89
|
||||
done, pending = cast(
|
||||
tuple[set[Awaitable[T]], set[Awaitable[T]]],
|
||||
await asyncio.wait(
|
||||
pending,
|
||||
timeout=remaining,
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
),
|
||||
)
|
||||
|
||||
for item in done:
|
||||
yield await item
|
||||
|
||||
|
||||
async def as_generated(
|
||||
iterables: Iterable[AsyncIterable[T]],
|
||||
*,
|
||||
return_exceptions: bool = False,
|
||||
) -> AsyncIterable[T]:
|
||||
"""
|
||||
Yield results from one or more async iterables, in the order they are produced.
|
||||
|
||||
Like :func:`as_completed`, but for async iterators or generators instead of futures.
|
||||
Creates a separate task to drain each iterable, and a single queue for results.
|
||||
|
||||
If ``return_exceptions`` is ``False``, then any exception will be raised, and
|
||||
pending iterables and tasks will be cancelled, and async generators will be closed.
|
||||
If ``return_exceptions`` is ``True``, any exceptions will be yielded as results,
|
||||
and execution will continue until all iterables have been fully consumed.
|
||||
|
||||
Example::
|
||||
|
||||
async def generator(x):
|
||||
for i in range(x):
|
||||
yield i
|
||||
|
||||
gen1 = generator(10)
|
||||
gen2 = generator(12)
|
||||
|
||||
async for value in as_generated([gen1, gen2]):
|
||||
... # intermixed values yielded from gen1 and gen2
|
||||
"""
|
||||
|
||||
exc_queue: asyncio.Queue[Exception] = asyncio.Queue()
|
||||
queue: asyncio.Queue[T] = asyncio.Queue()
|
||||
|
||||
async def tailer(iter: AsyncIterable[T]) -> None:
|
||||
try:
|
||||
async for item in iter:
|
||||
await queue.put(item)
|
||||
except asyncio.CancelledError:
|
||||
if isinstance(iter, AsyncGenerator): # pragma:nocover
|
||||
await iter.aclose()
|
||||
raise
|
||||
except Exception as e:
|
||||
await exc_queue.put(e)
|
||||
|
||||
tasks = [asyncio.ensure_future(tailer(iter)) for iter in iterables]
|
||||
pending = set(tasks)
|
||||
|
||||
try:
|
||||
while pending:
|
||||
try:
|
||||
exc = exc_queue.get_nowait()
|
||||
if return_exceptions:
|
||||
yield exc # type: ignore
|
||||
else:
|
||||
raise exc
|
||||
except asyncio.QueueEmpty:
|
||||
pass
|
||||
|
||||
try:
|
||||
value = queue.get_nowait()
|
||||
yield value
|
||||
except asyncio.QueueEmpty:
|
||||
for task in list(pending):
|
||||
if task.done():
|
||||
pending.remove(task)
|
||||
await asyncio.sleep(0.001)
|
||||
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
pass
|
||||
|
||||
finally:
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
|
||||
for task in tasks:
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
|
||||
async def gather(
|
||||
*args: Awaitable[T],
|
||||
return_exceptions: bool = False,
|
||||
limit: int = -1,
|
||||
) -> list[Any]:
|
||||
"""
|
||||
Like asyncio.gather but with a limit on concurrency.
|
||||
|
||||
Note that all results are buffered.
|
||||
|
||||
If gather is cancelled all tasks that were internally created and still pending
|
||||
will be cancelled as well.
|
||||
|
||||
Example::
|
||||
|
||||
futures = [some_coro(i) for i in range(10)]
|
||||
|
||||
results = await gather(*futures, limit=2)
|
||||
"""
|
||||
|
||||
# For detecting input duplicates and reconciling them at the end
|
||||
input_map: dict[Awaitable[T], list[int]] = {}
|
||||
# This is keyed on what we'll get back from asyncio.wait
|
||||
pos: dict[asyncio.Future[T], int] = {}
|
||||
ret: list[Any] = [None] * len(args)
|
||||
|
||||
pending: set[asyncio.Future[T]] = set()
|
||||
done: set[asyncio.Future[T]] = set()
|
||||
|
||||
next_arg = 0
|
||||
|
||||
while True:
|
||||
while next_arg < len(args) and (limit == -1 or len(pending) < limit):
|
||||
# We have to defer the creation of the Task as long as possible
|
||||
# because once we do, it starts executing, regardless of what we
|
||||
# have in the pending set.
|
||||
if args[next_arg] in input_map:
|
||||
input_map[args[next_arg]].append(next_arg)
|
||||
else:
|
||||
# We call ensure_future directly to ensure that we have a Task
|
||||
# because the return value of asyncio.wait will be an implicit
|
||||
# task otherwise, and we won't be able to know which input it
|
||||
# corresponds to.
|
||||
task: asyncio.Future[T] = asyncio.ensure_future(args[next_arg])
|
||||
pending.add(task)
|
||||
pos[task] = next_arg
|
||||
input_map[args[next_arg]] = [next_arg]
|
||||
next_arg += 1
|
||||
|
||||
# pending might be empty if the last items of args were dupes;
|
||||
# asyncio.wait([]) will raise an exception.
|
||||
if pending:
|
||||
try:
|
||||
done, pending = await asyncio.wait(
|
||||
pending, return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
for x in done:
|
||||
if return_exceptions and x.exception():
|
||||
ret[pos[x]] = x.exception()
|
||||
else:
|
||||
ret[pos[x]] = x.result()
|
||||
except asyncio.CancelledError:
|
||||
# Since we created these tasks we should cancel them
|
||||
for x in pending:
|
||||
x.cancel()
|
||||
# we insure that all tasks are cancelled before we raise
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
raise
|
||||
|
||||
if not pending and next_arg == len(args):
|
||||
break
|
||||
|
||||
for lst in input_map.values():
|
||||
for i in range(1, len(lst)):
|
||||
ret[lst[i]] = ret[lst[0]]
|
||||
|
||||
return ret
|
||||
|
||||
|
||||
async def gather_iter(
|
||||
itr: AnyIterable[MaybeAwaitable[T]],
|
||||
return_exceptions: bool = False,
|
||||
limit: int = -1,
|
||||
) -> list[T]:
|
||||
"""
|
||||
Wrapper around gather to handle gathering an iterable instead of ``*args``.
|
||||
|
||||
Note that the iterable values don't have to be awaitable.
|
||||
"""
|
||||
return await gather(
|
||||
*[maybe_await(i) async for i in aiter(itr)],
|
||||
return_exceptions=return_exceptions,
|
||||
limit=limit,
|
||||
)
|
||||
@@ -0,0 +1,427 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
"""
|
||||
Async-compatible versions of builtin functions for iterables.
|
||||
|
||||
These functions intentionally shadow their builtins counterparts,
|
||||
enabling use with both standard iterables and async iterables, without
|
||||
needing to use if/else clauses or awkward logic. Standard iterables
|
||||
get wrapped in async generators, and all functions are designed for
|
||||
use with `await`, `async for`, etc.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import builtins
|
||||
from collections.abc import AsyncIterable, AsyncIterator, Iterable
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, cast, Optional, overload, Union
|
||||
|
||||
from . import asyncio as ait_asyncio
|
||||
from .helpers import maybe_await, Orderable
|
||||
from .types import (
|
||||
AnyIterable,
|
||||
AnyIterator,
|
||||
AnyStop,
|
||||
MaybeAwaitable,
|
||||
R,
|
||||
T,
|
||||
T1,
|
||||
T2,
|
||||
T3,
|
||||
T4,
|
||||
T5,
|
||||
)
|
||||
|
||||
|
||||
class Sentinel(Enum):
|
||||
"""
|
||||
:meta private:
|
||||
"""
|
||||
|
||||
MISSING = object()
|
||||
|
||||
|
||||
async def all(itr: AnyIterable[MaybeAwaitable[Any]]) -> bool:
|
||||
"""
|
||||
Return True if all values are truthy in a mixed iterable, else False.
|
||||
The iterable will be fully consumed and any awaitables will
|
||||
automatically be awaited.
|
||||
|
||||
Example::
|
||||
|
||||
if await all(it):
|
||||
...
|
||||
|
||||
"""
|
||||
return builtins.all(await ait_asyncio.gather_iter(itr))
|
||||
|
||||
|
||||
async def any(itr: AnyIterable[MaybeAwaitable[Any]]) -> bool:
|
||||
"""
|
||||
Return True if any value is truthy in a mixed iterable, else False.
|
||||
The iterable will be fully consumed and any awaitables will
|
||||
automatically be awaited.
|
||||
|
||||
Example::
|
||||
|
||||
if await any(it):
|
||||
...
|
||||
|
||||
"""
|
||||
return builtins.any(await ait_asyncio.gather_iter(itr))
|
||||
|
||||
|
||||
def iter(itr: AnyIterable[T]) -> AsyncIterator[T]:
|
||||
"""
|
||||
Get an async iterator from any mixed iterable.
|
||||
|
||||
Async iterators will be returned directly.
|
||||
Async iterables will return an async iterator.
|
||||
Standard iterables will be wrapped in an async generator yielding
|
||||
each item in the iterable in the same order.
|
||||
|
||||
Examples::
|
||||
|
||||
async for value in iter(range(10)):
|
||||
...
|
||||
|
||||
"""
|
||||
if isinstance(itr, AsyncIterator):
|
||||
return itr
|
||||
|
||||
if isinstance(itr, AsyncIterable):
|
||||
return itr.__aiter__()
|
||||
|
||||
async def gen() -> AsyncIterator[T]:
|
||||
for item in cast(Iterable[T], itr):
|
||||
yield item
|
||||
|
||||
return gen()
|
||||
|
||||
|
||||
@overload
|
||||
async def next(itr: AnyIterator[T]) -> T: # pragma: no cover
|
||||
...
|
||||
|
||||
|
||||
@overload
|
||||
async def next(itr: AnyIterator[T1], default: T2) -> Union[T1, T2]: # pragma: no cover
|
||||
...
|
||||
|
||||
|
||||
async def next(
|
||||
itr: AnyIterator[T1], default: Union[T2, Sentinel] = Sentinel.MISSING
|
||||
) -> Union[T1, T2]:
|
||||
"""
|
||||
Return the next item of any mixed iterator.
|
||||
|
||||
Calls builtins.next() on standard iterators, and awaits itr.__anext__()
|
||||
on async iterators.
|
||||
|
||||
Example::
|
||||
|
||||
value = await next(it)
|
||||
|
||||
"""
|
||||
try:
|
||||
if isinstance(itr, AsyncIterator):
|
||||
return await itr.__anext__()
|
||||
|
||||
try:
|
||||
return builtins.next(itr)
|
||||
except StopIteration:
|
||||
raise StopAsyncIteration
|
||||
except StopAsyncIteration:
|
||||
if default is Sentinel.MISSING:
|
||||
raise
|
||||
return default
|
||||
|
||||
|
||||
async def list(itr: AnyIterable[T]) -> builtins.list[T]:
|
||||
"""
|
||||
Consume a mixed iterable and return a list of items in order.
|
||||
|
||||
Example::
|
||||
|
||||
await list(range(5))
|
||||
-> [0, 1, 2, 3, 4]
|
||||
|
||||
"""
|
||||
return [item async for item in iter(itr)]
|
||||
|
||||
|
||||
async def tuple(itr: AnyIterable[T]) -> builtins.tuple[T, ...]:
|
||||
"""
|
||||
Consume a mixed iterable and return a tuple of items in order.
|
||||
|
||||
Example::
|
||||
|
||||
await tuple(range(5))
|
||||
-> (0, 1, 2, 3, 4)
|
||||
|
||||
"""
|
||||
# Suboptimal, but tuple can't be created from AsyncIterable directly.
|
||||
return builtins.tuple(await list(itr))
|
||||
|
||||
|
||||
async def set(itr: AnyIterable[T]) -> builtins.set[T]:
|
||||
"""
|
||||
Consume a mixed iterable and return a set of items.
|
||||
|
||||
Example::
|
||||
|
||||
await set([0, 1, 2, 3, 0, 1, 2, 3])
|
||||
-> {0, 1, 2, 3}
|
||||
|
||||
"""
|
||||
return {item async for item in iter(itr)}
|
||||
|
||||
|
||||
async def enumerate(
|
||||
itr: AnyIterable[T], start: int = 0
|
||||
) -> AsyncIterator[builtins.tuple[int, T]]:
|
||||
"""
|
||||
Consume a mixed iterable and yield the current index and item.
|
||||
|
||||
Example::
|
||||
|
||||
async for index, value in enumerate(...):
|
||||
...
|
||||
|
||||
"""
|
||||
index = start
|
||||
async for item in iter(itr):
|
||||
yield index, item
|
||||
index += 1
|
||||
|
||||
|
||||
async def map(fn: Callable[[T], R], itr: AnyIterable[T]) -> AsyncIterator[R]:
|
||||
"""
|
||||
Modify item of a mixed iterable using the given function or coroutine.
|
||||
|
||||
Example::
|
||||
|
||||
async for response in map(func, data):
|
||||
...
|
||||
|
||||
"""
|
||||
# todo: queue items eagerly
|
||||
async for item in iter(itr):
|
||||
yield await maybe_await(fn(item))
|
||||
|
||||
|
||||
@overload
|
||||
async def max(
|
||||
itr: AnyIterable[Orderable], *, key: Optional[Callable] = None
|
||||
) -> Orderable: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
async def max(
|
||||
itr: AnyIterable[Orderable], *, default: T, key: Optional[Callable] = None
|
||||
) -> Union[Orderable, T]: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
async def max(itr: AnyIterable[Orderable], **kwargs: Any) -> Any:
|
||||
"""
|
||||
Return the largest item in an iterable or the largest of two or more arguments.
|
||||
|
||||
Example::
|
||||
|
||||
await max(range(5))
|
||||
-> 4
|
||||
|
||||
"""
|
||||
for k in kwargs:
|
||||
if k not in ("key", "default"):
|
||||
raise ValueError(f"kwarg {k} not supported")
|
||||
|
||||
value: Orderable
|
||||
vkey: Any
|
||||
|
||||
keyfunc = kwargs.get("key", None)
|
||||
it = iter(itr)
|
||||
|
||||
try:
|
||||
value = await next(it)
|
||||
if keyfunc:
|
||||
vkey = keyfunc(value)
|
||||
|
||||
except StopAsyncIteration:
|
||||
if "default" in kwargs:
|
||||
return kwargs["default"]
|
||||
raise ValueError("iterable is empty and no default value given")
|
||||
|
||||
if keyfunc:
|
||||
async for item in it:
|
||||
ikey = keyfunc(item)
|
||||
if ikey > vkey:
|
||||
value = item
|
||||
vkey = ikey
|
||||
|
||||
else:
|
||||
async for item in it:
|
||||
if item > value:
|
||||
value = item
|
||||
|
||||
return value
|
||||
|
||||
|
||||
@overload
|
||||
async def min(
|
||||
itr: AnyIterable[Orderable], *, key: Optional[Callable] = None
|
||||
) -> Orderable: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
async def min(
|
||||
itr: AnyIterable[Orderable], *, default: T, key: Optional[Callable] = None
|
||||
) -> Union[Orderable, T]: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
async def min(itr: AnyIterable[Orderable], **kwargs: Any) -> Any:
|
||||
"""
|
||||
Return the smallest item in an iterable or the smallest of two or more arguments.
|
||||
|
||||
Example::
|
||||
|
||||
await min(range(5))
|
||||
-> 0
|
||||
|
||||
"""
|
||||
for k in kwargs:
|
||||
if k not in ("key", "default"):
|
||||
raise ValueError(f"kwarg {k} not supported")
|
||||
|
||||
value: Orderable
|
||||
vkey: Any
|
||||
|
||||
keyfunc = kwargs.get("key", None)
|
||||
it = iter(itr)
|
||||
|
||||
try:
|
||||
value = await next(it)
|
||||
if keyfunc:
|
||||
vkey = keyfunc(value)
|
||||
|
||||
except StopAsyncIteration:
|
||||
if "default" in kwargs:
|
||||
return kwargs["default"]
|
||||
raise ValueError("iterable is empty and no default value given")
|
||||
|
||||
if keyfunc:
|
||||
async for item in it:
|
||||
ikey = keyfunc(item)
|
||||
if ikey < vkey:
|
||||
value = item
|
||||
vkey = ikey
|
||||
|
||||
else:
|
||||
async for item in it:
|
||||
if item < value:
|
||||
value = item
|
||||
|
||||
return value
|
||||
|
||||
|
||||
async def sum(itr: AnyIterable[T], start: Optional[T] = None) -> T:
|
||||
"""
|
||||
Compute the sum of a mixed iterable, adding each value with the start value.
|
||||
|
||||
Example::
|
||||
|
||||
await sum(generator())
|
||||
-> 1024
|
||||
|
||||
"""
|
||||
value: T
|
||||
if start is None:
|
||||
value = cast(T, 0) # emulate stdlib but still type nicely for non-ints
|
||||
else:
|
||||
value = start
|
||||
|
||||
async for item in iter(itr):
|
||||
value += item # type: ignore # mypy doesn't know T + T
|
||||
|
||||
return value
|
||||
|
||||
|
||||
@overload
|
||||
def zip(
|
||||
__iter1: AnyIterable[T1],
|
||||
) -> AsyncIterator[builtins.tuple[T1]]: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
def zip(
|
||||
__iter1: AnyIterable[T1], __iter2: AnyIterable[T2]
|
||||
) -> AsyncIterator[builtins.tuple[T1, T2]]: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
def zip(
|
||||
__iter1: AnyIterable[T1], __iter2: AnyIterable[T2], __iter3: AnyIterable[T3]
|
||||
) -> AsyncIterator[builtins.tuple[T1, T2, T3]]: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
def zip(
|
||||
__iter1: AnyIterable[T1],
|
||||
__iter2: AnyIterable[T2],
|
||||
__iter3: AnyIterable[T3],
|
||||
__iter4: AnyIterable[T4],
|
||||
) -> AsyncIterator[builtins.tuple[T1, T2, T3, T4]]: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
def zip(
|
||||
__iter1: AnyIterable[T1],
|
||||
__iter2: AnyIterable[T2],
|
||||
__iter3: AnyIterable[T3],
|
||||
__iter4: AnyIterable[T4],
|
||||
__iter5: AnyIterable[T5],
|
||||
) -> AsyncIterator[builtins.tuple[T1, T2, T3, T4, T5]]: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
def zip(
|
||||
__iter1: AnyIterable[Any],
|
||||
__iter2: AnyIterable[Any],
|
||||
__iter3: AnyIterable[Any],
|
||||
__iter4: AnyIterable[Any],
|
||||
__iter5: AnyIterable[Any],
|
||||
__iter6: AnyIterable[Any],
|
||||
*__iterables: AnyIterable[Any],
|
||||
) -> AsyncIterator[builtins.tuple[Any, ...]]: # pragma: no cover
|
||||
pass
|
||||
|
||||
|
||||
async def zip(*itrs: AnyIterable[Any]) -> AsyncIterator[builtins.tuple[Any, ...]]:
|
||||
"""
|
||||
Yield a tuple of items from mixed iterables until the shortest is consumed.
|
||||
|
||||
Example::
|
||||
|
||||
async for a, b, c in zip(i, j, k):
|
||||
...
|
||||
|
||||
"""
|
||||
its: builtins.list[AsyncIterator[Any]] = [iter(itr) for itr in itrs]
|
||||
|
||||
while True:
|
||||
values = await asyncio.gather(
|
||||
*[it.__anext__() for it in its], return_exceptions=True
|
||||
)
|
||||
if builtins.any(isinstance(v, AnyStop) for v in values):
|
||||
break
|
||||
yield builtins.tuple(values)
|
||||
@@ -0,0 +1,21 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
import inspect
|
||||
from collections.abc import Awaitable
|
||||
|
||||
from typing import Protocol, Union
|
||||
|
||||
from .types import T
|
||||
|
||||
|
||||
class Orderable(Protocol): # pragma: no cover
|
||||
def __lt__(self, other): ...
|
||||
|
||||
def __gt__(self, other): ...
|
||||
|
||||
|
||||
async def maybe_await(object: Union[Awaitable[T], T]) -> T:
|
||||
if inspect.isawaitable(object):
|
||||
return await object # type: ignore
|
||||
return object # type: ignore
|
||||
@@ -0,0 +1,589 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
"""
|
||||
Async-compatible version of itertools standard library functions.
|
||||
|
||||
These functions build on top of the async builtins components,
|
||||
enabling use of both standard iterables and async iterables, without
|
||||
needing to use if/else clauses or awkward logic. Standard iterables
|
||||
get wrapped in async generators, and all functions are designed for
|
||||
use with `await`, `async for`, etc.
|
||||
|
||||
See https://docs.python.org/3/library/itertools.html for reference.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import builtins
|
||||
import itertools
|
||||
import operator
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, Optional, overload
|
||||
|
||||
from .builtins import enumerate, iter, list, next, tuple, zip
|
||||
from .helpers import maybe_await
|
||||
from .types import (
|
||||
Accumulator,
|
||||
AnyFunction,
|
||||
AnyIterable,
|
||||
AnyIterableIterable,
|
||||
AnyStop,
|
||||
KeyFunction,
|
||||
N,
|
||||
Predicate,
|
||||
R,
|
||||
T,
|
||||
)
|
||||
|
||||
|
||||
async def accumulate(
|
||||
itr: AnyIterable[T], func: Accumulator[T] = operator.add
|
||||
) -> AsyncIterator[T]:
|
||||
"""
|
||||
Yield the running accumulation of an iterable and operator.
|
||||
|
||||
Accepts both a standard function or a coroutine for accumulation.
|
||||
|
||||
Example::
|
||||
|
||||
data = [1, 2, 3, 4]
|
||||
|
||||
async def mul(a, b):
|
||||
return a * b
|
||||
|
||||
async for total in accumulate(data, func=mul):
|
||||
... # 1, 2, 6, 24
|
||||
|
||||
"""
|
||||
itr = iter(itr)
|
||||
try:
|
||||
total: T = await next(itr)
|
||||
except AnyStop:
|
||||
return
|
||||
|
||||
yield total
|
||||
async for item in itr:
|
||||
total = await maybe_await(func(total, item))
|
||||
yield total
|
||||
|
||||
|
||||
async def batched(
|
||||
iterable: AnyIterable[T],
|
||||
n: int,
|
||||
*,
|
||||
strict: bool = False,
|
||||
) -> AsyncIterator[builtins.tuple[T, ...]]:
|
||||
"""
|
||||
Yield batches of values from the given iterable. The final batch may be shorter.
|
||||
|
||||
Example::
|
||||
|
||||
async for batch in batched(range(15), 5):
|
||||
... # (0, 1, 2, 3, 4), (5, 6, 7, 8, 9), (10, 11, 12, 13, 14)
|
||||
|
||||
"""
|
||||
if n < 1:
|
||||
raise ValueError("n must be at least one")
|
||||
aiterator = iter(iterable)
|
||||
while batch := await tuple(islice(aiterator, n)):
|
||||
if strict and len(batch) != n:
|
||||
raise ValueError("batched: incomplete batch")
|
||||
yield batch
|
||||
|
||||
|
||||
class Chain:
|
||||
def __call__(self, *itrs: AnyIterable[T]) -> AsyncIterator[T]:
|
||||
"""
|
||||
Yield values from one or more iterables in series.
|
||||
|
||||
Consumes the first iterable lazily, in entirety, then the second, and so on.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in chain([1, 2, 3], [7, 8, 9]):
|
||||
... # 1, 2, 3, 7, 8, 9
|
||||
|
||||
"""
|
||||
return self.from_iterable(itrs)
|
||||
|
||||
async def from_iterable(self, itrs: AnyIterableIterable[T]) -> AsyncIterator[T]:
|
||||
"""
|
||||
Like chain, but takes an iterable of iterables.
|
||||
|
||||
Alias for chain(*itrs)
|
||||
"""
|
||||
async for itr in iter(itrs):
|
||||
async for item in iter(itr):
|
||||
yield item
|
||||
|
||||
|
||||
chain = Chain()
|
||||
|
||||
|
||||
async def combinations(
|
||||
itr: AnyIterable[T], r: int
|
||||
) -> AsyncIterator[builtins.tuple[T, ...]]:
|
||||
"""
|
||||
Yield r length subsequences from the given iterable.
|
||||
|
||||
Simple wrapper around itertools.combinations for asyncio.
|
||||
This will consume the entire iterable before yielding values.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in combinations(range(4), 3):
|
||||
... # (0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3)
|
||||
|
||||
"""
|
||||
pool: builtins.list[T] = await list(itr)
|
||||
for value in itertools.combinations(pool, r):
|
||||
yield value
|
||||
|
||||
|
||||
async def combinations_with_replacement(
|
||||
itr: AnyIterable[T], r: int
|
||||
) -> AsyncIterator[builtins.tuple[T, ...]]:
|
||||
"""
|
||||
Yield r length subsequences from the given iterable with replacement.
|
||||
|
||||
Simple wrapper around itertools.combinations_with_replacement.
|
||||
This will consume the entire iterable before yielding values.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in combinations_with_replacement("ABC", 2):
|
||||
... # ("A", "A"), ("A", "B"), ("A", "C"), ("B", "B"), ...
|
||||
|
||||
"""
|
||||
pool: builtins.list[T] = await list(itr)
|
||||
for value in itertools.combinations_with_replacement(pool, r):
|
||||
yield value
|
||||
|
||||
|
||||
async def compress(
|
||||
itr: AnyIterable[T], selectors: AnyIterable[Any]
|
||||
) -> AsyncIterator[T]:
|
||||
"""
|
||||
Yield elements only when the corresponding selector evaluates to True.
|
||||
|
||||
Stops when either the iterable or the selectors have been exhausted.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in compress(range(5), [1, 0, 0, 1, 1]):
|
||||
... # 0, 3, 4
|
||||
"""
|
||||
async for value, selector in zip(itr, selectors):
|
||||
if selector:
|
||||
yield value
|
||||
|
||||
|
||||
async def count(start: N = 0, step: N = 1) -> AsyncIterator[N]:
|
||||
"""
|
||||
Yield an infinite series, starting at the given value and increasing by step.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in counter(10, -1):
|
||||
... # 10, 9, 8, 7, ...
|
||||
|
||||
"""
|
||||
|
||||
value = start
|
||||
while True:
|
||||
yield value
|
||||
value += step
|
||||
|
||||
|
||||
async def cycle(itr: AnyIterable[T]) -> AsyncIterator[T]:
|
||||
"""
|
||||
Yield a repeating series from the given iterable.
|
||||
|
||||
Lazily consumes the iterable when the next value is needed, and caching
|
||||
the values in memory for future iterations of the series.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in cycle([1, 2]):
|
||||
... # 1, 2, 1, 2, 1, 2, ...
|
||||
|
||||
"""
|
||||
items = []
|
||||
async for item in iter(itr):
|
||||
yield item
|
||||
items.append(item)
|
||||
|
||||
while True:
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
|
||||
async def dropwhile(
|
||||
predicate: Predicate[T], iterable: AnyIterable[T]
|
||||
) -> AsyncIterator[T]:
|
||||
"""
|
||||
Drops all items until the predicate evaluates False; yields all items afterwards.
|
||||
|
||||
Accepts both standard functions and coroutines for the predicate.
|
||||
|
||||
Example::
|
||||
|
||||
def pred(x):
|
||||
return x < 4
|
||||
|
||||
async for item in dropwhile(pred, range(6)):
|
||||
... # 4, 5, 6
|
||||
|
||||
"""
|
||||
itr = iter(iterable)
|
||||
async for item in itr:
|
||||
if not await maybe_await(predicate(item)):
|
||||
yield item
|
||||
break
|
||||
async for item in itr:
|
||||
yield item
|
||||
|
||||
|
||||
async def filterfalse(
|
||||
predicate: Predicate[T], iterable: AnyIterable[T]
|
||||
) -> AsyncIterator[T]:
|
||||
"""
|
||||
Yield items from the iterable only when the predicate evaluates to False.
|
||||
|
||||
Accepts both standard functions and coroutines for the predicate.
|
||||
|
||||
Example::
|
||||
|
||||
def pred(x):
|
||||
return x < 4
|
||||
|
||||
async for item in filterfalse(pred, range(6)):
|
||||
... # 4, 5
|
||||
|
||||
"""
|
||||
async for item in iter(iterable):
|
||||
if not await maybe_await(predicate(item)):
|
||||
yield item
|
||||
|
||||
|
||||
@overload
|
||||
def groupby(
|
||||
itr: AnyIterable[T],
|
||||
) -> AsyncIterator[builtins.tuple[T, builtins.list[T]]]: # pragma: nocover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
def groupby(
|
||||
itr: AnyIterable[T], key: KeyFunction[T, R]
|
||||
) -> AsyncIterator[builtins.tuple[R, builtins.list[T]]]: # pragma: nocover
|
||||
pass
|
||||
|
||||
|
||||
async def groupby(
|
||||
itr: AnyIterable[T], key: Optional[KeyFunction[T, R]] = None
|
||||
) -> AsyncIterator[builtins.tuple[Any, builtins.list[T]]]:
|
||||
"""
|
||||
Yield consecutive keys and groupings from the given iterable.
|
||||
|
||||
Items will be grouped based on the key function, which defaults to
|
||||
the identity of each item. Accepts both standard functions and
|
||||
coroutines for the key function. Suggest sorting by the key
|
||||
function before using groupby.
|
||||
|
||||
Example::
|
||||
|
||||
data = ["A", "a", "b", "c", "C", "c"]
|
||||
|
||||
async for key, group in groupby(data, key=str.lower):
|
||||
key # "a", "b", "c"
|
||||
group # ["A", "a"], ["b"], ["c", "C", "c"]
|
||||
|
||||
"""
|
||||
if key is None:
|
||||
key = lambda x: x # noqa: E731
|
||||
|
||||
grouping: builtins.list[T] = []
|
||||
|
||||
it = iter(itr)
|
||||
try:
|
||||
item = await next(it)
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
grouping = [item]
|
||||
|
||||
j = await maybe_await(key(item))
|
||||
async for item in it:
|
||||
k = await maybe_await(key(item))
|
||||
if k != j:
|
||||
yield j, grouping
|
||||
grouping = [item]
|
||||
else:
|
||||
grouping.append(item)
|
||||
j = k
|
||||
|
||||
yield j, grouping
|
||||
|
||||
|
||||
@overload
|
||||
def islice(
|
||||
itr: AnyIterable[T], __stop: Optional[int]
|
||||
) -> AsyncIterator[T]: # pragma: nocover
|
||||
pass
|
||||
|
||||
|
||||
@overload
|
||||
def islice(
|
||||
itr: AnyIterable[T], __start: int, __stop: Optional[int], __step: int = 1
|
||||
) -> AsyncIterator[T]: # pragma: nocover
|
||||
pass
|
||||
|
||||
|
||||
async def islice(itr: AnyIterable[T], *args: Optional[int]) -> AsyncIterator[T]:
|
||||
"""
|
||||
Yield selected items from the given iterable.
|
||||
|
||||
islice(iterable, stop)
|
||||
islice(iterable, start, stop[, step])
|
||||
|
||||
Starting from the start index (or zero), stopping at the stop
|
||||
index (or until exhausted), skipping items if step > 0.
|
||||
|
||||
Example::
|
||||
|
||||
data = range(10)
|
||||
|
||||
async for item in islice(data, 5):
|
||||
... # 0, 1, 2, 3, 4
|
||||
|
||||
async for item in islice(data, 2, 5):
|
||||
... # 2, 3, 4
|
||||
|
||||
async for item in islice(data, 1, 7, 2):
|
||||
... # 1, 3, 5
|
||||
|
||||
"""
|
||||
start = 0
|
||||
step = 1
|
||||
if not args:
|
||||
raise ValueError("must pass stop index")
|
||||
if len(args) == 1:
|
||||
(stop,) = args
|
||||
elif len(args) == 2:
|
||||
start, stop = args # type: ignore
|
||||
elif len(args) == 3:
|
||||
start, stop, step = args # type: ignore
|
||||
else:
|
||||
raise ValueError("too many arguments given")
|
||||
assert start >= 0 and (stop is None or stop >= 0) and step >= 0
|
||||
step = max(1, step)
|
||||
|
||||
if stop == 0:
|
||||
return
|
||||
|
||||
async for index, item in enumerate(itr):
|
||||
if index >= start and (index - start) % step == 0:
|
||||
yield item
|
||||
if stop is not None and index + 1 >= stop:
|
||||
break
|
||||
|
||||
|
||||
async def permutations(
|
||||
itr: AnyIterable[T], r: Optional[int] = None
|
||||
) -> AsyncIterator[builtins.tuple[T, ...]]:
|
||||
"""
|
||||
Yield r length permutations of elements in the iterable.
|
||||
|
||||
Simple wrapper around itertools.combinations for asyncio.
|
||||
This will consume the entire iterable before yielding values.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in permutations(range(3)):
|
||||
... # (0, 1, 2), (0, 2, 1), (1, 0, 2), ...
|
||||
|
||||
"""
|
||||
pool: builtins.list[T] = await list(itr)
|
||||
for value in itertools.permutations(pool, r):
|
||||
yield value
|
||||
|
||||
|
||||
async def product(
|
||||
*itrs: AnyIterable[T], repeat: int = 1
|
||||
) -> AsyncIterator[builtins.tuple[T, ...]]:
|
||||
"""
|
||||
Yield cartesian products of all iterables.
|
||||
|
||||
Simple wrapper around itertools.combinations for asyncio.
|
||||
This will consume all iterables before yielding any values.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in product("abc", "xy"):
|
||||
... # ("a", "x"), ("a", "y"), ("b", "x"), ...
|
||||
|
||||
async for value in product(range(3), repeat=3):
|
||||
... # (0, 0, 0), (0, 0, 1), (0, 0, 2), ...
|
||||
|
||||
"""
|
||||
pools = await asyncio.gather(*[list(itr) for itr in itrs])
|
||||
value: builtins.tuple[T, ...]
|
||||
for value in itertools.product(*pools, repeat=repeat):
|
||||
yield value
|
||||
|
||||
|
||||
async def repeat(elem: T, n: int = -1) -> AsyncIterator[T]:
|
||||
"""
|
||||
Yield the given value repeatedly, forever or up to n times.
|
||||
|
||||
Example::
|
||||
|
||||
async for value in repeat(7):
|
||||
... # 7, 7, 7, 7, 7, 7, ...
|
||||
|
||||
"""
|
||||
while True:
|
||||
if n == 0:
|
||||
break
|
||||
yield elem
|
||||
n -= 1
|
||||
|
||||
|
||||
async def starmap(
|
||||
fn: AnyFunction[R], iterable: AnyIterableIterable[Any]
|
||||
) -> AsyncIterator[R]:
|
||||
"""
|
||||
Yield values from a function using an iterable of iterables for arguments.
|
||||
|
||||
Each iterable contained within will be unpacked and consumed before
|
||||
executing the function or coroutine.
|
||||
|
||||
Example::
|
||||
|
||||
data = [(1, 1), (1, 1, 1), (2, 2)]
|
||||
|
||||
async for value in starmap(operator.add, data):
|
||||
... # 2, 3, 4
|
||||
|
||||
"""
|
||||
async for itr in iter(iterable):
|
||||
args = await list(itr)
|
||||
yield await maybe_await(fn(*args))
|
||||
|
||||
|
||||
async def takewhile(
|
||||
predicate: Predicate[T], iterable: AnyIterable[T]
|
||||
) -> AsyncIterator[T]:
|
||||
"""
|
||||
Yield values from the iterable until the predicate evaluates False.
|
||||
|
||||
Accepts both standard functions and coroutines for the predicate.
|
||||
|
||||
Example::
|
||||
|
||||
def pred(x):
|
||||
return x < 4
|
||||
|
||||
async for value in takewhile(pred, range(8)):
|
||||
... # 0, 1, 2, 3
|
||||
|
||||
"""
|
||||
async for item in iter(iterable):
|
||||
if await maybe_await(predicate(item)):
|
||||
yield item
|
||||
else:
|
||||
break
|
||||
|
||||
|
||||
def tee(itr: AnyIterable[T], n: int = 2) -> builtins.tuple[AsyncIterator[T], ...]:
|
||||
"""
|
||||
Return n iterators that each yield items from the given iterable.
|
||||
|
||||
The first iterator lazily fetches from the original iterable, and then
|
||||
queues the values for the other iterators to yield when needed.
|
||||
|
||||
Caveat: all iterators are dependent on the first iterator — if it is
|
||||
consumed more slowly than the rest, the other consumers will be blocked
|
||||
until the first iterator continues forward. Similarly, if the first
|
||||
iterator is consumed more quickly than the rest, more memory will be
|
||||
used in keeping values in the queues until the other iterators finish
|
||||
consuming them.
|
||||
|
||||
Example::
|
||||
|
||||
it1, it2 = tee(range(5), n=2)
|
||||
|
||||
async for value in it1:
|
||||
... # 0, 1, 2, 3, 4
|
||||
|
||||
async for value in it2:
|
||||
... # 0, 1, 2, 3, 4
|
||||
|
||||
"""
|
||||
assert n > 0
|
||||
sentinel = object()
|
||||
queues: builtins.list[asyncio.Queue] = [asyncio.Queue() for k in range(n)]
|
||||
|
||||
async def gen(k: int, q: asyncio.Queue) -> AsyncIterator[T]:
|
||||
if k == 0:
|
||||
try:
|
||||
async for value in iter(itr):
|
||||
await asyncio.gather(
|
||||
*[queue.put((None, value)) for queue in queues[1:]]
|
||||
)
|
||||
yield value
|
||||
except Exception as e:
|
||||
await asyncio.gather(*[queue.put((e, None)) for queue in queues[1:]])
|
||||
raise
|
||||
|
||||
await asyncio.gather(*[queue.put((None, sentinel)) for queue in queues[1:]])
|
||||
|
||||
else:
|
||||
while True:
|
||||
error, value = await q.get()
|
||||
if error is not None:
|
||||
raise error
|
||||
if value is sentinel:
|
||||
break
|
||||
yield value
|
||||
|
||||
return builtins.tuple(gen(k, q) for k, q in builtins.enumerate(queues))
|
||||
|
||||
|
||||
async def zip_longest(
|
||||
*itrs: AnyIterable[Any], fillvalue: Any = None
|
||||
) -> AsyncIterator[builtins.tuple[Any, ...]]:
|
||||
"""
|
||||
Yield a tuple of items from mixed iterables until all are consumed.
|
||||
|
||||
If shorter iterables are exhausted, the default value will be used
|
||||
until all iterables are exhausted.
|
||||
|
||||
Example::
|
||||
|
||||
a = range(3)
|
||||
b = range(5)
|
||||
|
||||
async for a, b in zip_longest(a, b, fillvalue=-1):
|
||||
a # 0, 1, 2, -1, -1
|
||||
b # 0, 1, 2, 3, 4
|
||||
|
||||
"""
|
||||
its: builtins.list[AsyncIterator[Any]] = [iter(itr) for itr in itrs]
|
||||
itr_count = len(its)
|
||||
finished = 0
|
||||
|
||||
while True:
|
||||
values = await asyncio.gather(
|
||||
*[it.__anext__() for it in its], return_exceptions=True
|
||||
)
|
||||
for idx, value in builtins.enumerate(values):
|
||||
if isinstance(value, AnyStop):
|
||||
finished += 1
|
||||
values[idx] = fillvalue
|
||||
its[idx] = repeat(fillvalue)
|
||||
elif isinstance(value, BaseException):
|
||||
raise value
|
||||
if finished >= itr_count:
|
||||
break
|
||||
yield builtins.tuple(values)
|
||||
@@ -0,0 +1,95 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterable
|
||||
from typing import TypeVar
|
||||
|
||||
from aioitertools.helpers import maybe_await
|
||||
|
||||
from .builtins import iter
|
||||
from .itertools import islice
|
||||
from .types import AnyIterable, Predicate
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
async def take(n: int, iterable: AnyIterable[T]) -> list[T]:
|
||||
"""
|
||||
Return the first n items of iterable as a list.
|
||||
|
||||
If there are too few items in iterable, all of them are returned.
|
||||
n needs to be at least 0. If it is 0, an empty list is returned.
|
||||
|
||||
Example::
|
||||
|
||||
first_two = await take(2, [1, 2, 3, 4, 5])
|
||||
|
||||
"""
|
||||
if n < 0:
|
||||
raise ValueError("take's first parameter can't be negative")
|
||||
return [item async for item in islice(iterable, n)]
|
||||
|
||||
|
||||
async def chunked(iterable: AnyIterable[T], n: int) -> AsyncIterable[list[T]]:
|
||||
"""
|
||||
Break iterable into chunks of length n.
|
||||
|
||||
The last chunk will be shorter if the total number of items is not
|
||||
divisible by n.
|
||||
|
||||
Example::
|
||||
|
||||
async for chunk in chunked([1, 2, 3, 4, 5], n=2):
|
||||
... # first iteration: chunk == [1, 2]; last one: chunk == [5]
|
||||
"""
|
||||
it = iter(iterable)
|
||||
chunk = await take(n, it)
|
||||
while chunk != []:
|
||||
yield chunk
|
||||
chunk = await take(n, it)
|
||||
|
||||
|
||||
async def before_and_after(
|
||||
predicate: Predicate[T], iterable: AnyIterable[T]
|
||||
) -> tuple[AsyncIterable[T], AsyncIterable[T]]:
|
||||
"""
|
||||
A variant of :func:`aioitertools.takewhile` that allows complete access to the
|
||||
remainder of the iterator.
|
||||
|
||||
>>> it = iter('ABCdEfGhI')
|
||||
>>> all_upper, remainder = await before_and_after(str.isupper, it)
|
||||
>>> ''.join([char async for char in all_upper])
|
||||
'ABC'
|
||||
>>> ''.join([char async for char in remainder])
|
||||
'dEfGhI'
|
||||
|
||||
Note that the first iterator must be fully consumed before the second
|
||||
iterator can generate valid results.
|
||||
"""
|
||||
|
||||
it = iter(iterable)
|
||||
|
||||
transition = asyncio.get_event_loop().create_future()
|
||||
|
||||
async def true_iterator():
|
||||
async for elem in it:
|
||||
if await maybe_await(predicate(elem)):
|
||||
yield elem
|
||||
else:
|
||||
transition.set_result(elem)
|
||||
return
|
||||
|
||||
transition.set_exception(StopAsyncIteration)
|
||||
|
||||
async def remainder_iterator():
|
||||
try:
|
||||
yield (await transition)
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
|
||||
async for elm in it:
|
||||
yield elm
|
||||
|
||||
return true_iterator(), remainder_iterator()
|
||||
@@ -0,0 +1,8 @@
|
||||
# Copyright Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
from .asyncio import AsyncioTest
|
||||
from .builtins import BuiltinsTest
|
||||
from .helpers import HelpersTest
|
||||
from .itertools import ItertoolsTest
|
||||
from .more_itertools import MoreItertoolsTest
|
||||
@@ -0,0 +1,7 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
import unittest
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
unittest.main(module="aioitertools.tests", verbosity=2)
|
||||
@@ -0,0 +1,259 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
import asyncio
|
||||
from unittest import TestCase
|
||||
|
||||
import aioitertools as ait
|
||||
import aioitertools.asyncio as aio
|
||||
from .helpers import async_test
|
||||
|
||||
slist = ["A", "B", "C"]
|
||||
srange = range(3)
|
||||
|
||||
|
||||
class AsyncioTest(TestCase):
|
||||
def test_import(self):
|
||||
self.assertEqual(ait.asyncio, aio)
|
||||
|
||||
@async_test
|
||||
async def test_as_completed(self):
|
||||
async def sleepy(number, duration):
|
||||
await asyncio.sleep(duration)
|
||||
return number
|
||||
|
||||
pairs = [(1, 0.3), (2, 0.1), (3, 0.5)]
|
||||
expected = [2, 1, 3]
|
||||
|
||||
futures = [sleepy(*pair) for pair in pairs]
|
||||
results = await ait.list(aio.as_completed(futures))
|
||||
self.assertEqual(results, expected)
|
||||
|
||||
futures = [sleepy(*pair) for pair in pairs]
|
||||
results = []
|
||||
async for value in aio.as_completed(futures):
|
||||
results.append(value)
|
||||
self.assertEqual(results, expected)
|
||||
|
||||
@async_test
|
||||
async def test_as_completed_timeout(self):
|
||||
calls = [(1.0,), (0.1,)]
|
||||
|
||||
futures = [asyncio.sleep(*args) for args in calls]
|
||||
with self.assertRaises(asyncio.TimeoutError):
|
||||
await ait.list(aio.as_completed(futures, timeout=0.5))
|
||||
|
||||
futures = [asyncio.sleep(*args) for args in calls]
|
||||
results = 0
|
||||
with self.assertRaises(asyncio.TimeoutError):
|
||||
async for _ in aio.as_completed(futures, timeout=0.5):
|
||||
results += 1
|
||||
self.assertEqual(results, 1)
|
||||
|
||||
@async_test
|
||||
async def test_as_generated(self):
|
||||
async def gen():
|
||||
for i in range(10):
|
||||
yield i
|
||||
await asyncio.sleep(0)
|
||||
|
||||
gens = [gen(), gen(), gen()]
|
||||
expected = list(range(10)) * 3
|
||||
results = []
|
||||
async for value in aio.as_generated(gens):
|
||||
results.append(value)
|
||||
self.assertEqual(30, len(results))
|
||||
self.assertListEqual(sorted(expected), sorted(results))
|
||||
|
||||
@async_test
|
||||
async def test_as_generated_exception(self):
|
||||
async def gen1():
|
||||
for i in range(3):
|
||||
yield i
|
||||
await asyncio.sleep(0)
|
||||
raise Exception("fake")
|
||||
|
||||
async def gen2():
|
||||
for i in range(10):
|
||||
yield i
|
||||
await asyncio.sleep(0)
|
||||
|
||||
gens = [gen1(), gen2()]
|
||||
results = []
|
||||
with self.assertRaisesRegex(Exception, "fake"):
|
||||
async for value in aio.as_generated(gens):
|
||||
results.append(value)
|
||||
self.assertNotIn(10, results)
|
||||
|
||||
@async_test
|
||||
async def test_as_generated_return_exception(self):
|
||||
async def gen1():
|
||||
for i in range(3):
|
||||
yield i
|
||||
await asyncio.sleep(0)
|
||||
raise Exception("fake")
|
||||
|
||||
async def gen2():
|
||||
for i in range(10):
|
||||
yield i
|
||||
await asyncio.sleep(0)
|
||||
|
||||
gens = [gen1(), gen2()]
|
||||
expected = list(range(3)) + list(range(10))
|
||||
errors = []
|
||||
results = []
|
||||
async for value in aio.as_generated(gens, return_exceptions=True):
|
||||
if isinstance(value, Exception):
|
||||
errors.append(value)
|
||||
else:
|
||||
results.append(value)
|
||||
self.assertListEqual(sorted(expected), sorted(results))
|
||||
self.assertEqual(1, len(errors))
|
||||
self.assertIsInstance(errors[0], Exception)
|
||||
|
||||
@async_test
|
||||
async def test_as_generated_task_cancelled(self):
|
||||
async def gen(max: int = 10):
|
||||
for i in range(5):
|
||||
if i > max:
|
||||
raise asyncio.CancelledError
|
||||
yield i
|
||||
await asyncio.sleep(0)
|
||||
|
||||
gens = [gen(2), gen()]
|
||||
expected = list(range(3)) + list(range(5))
|
||||
results = []
|
||||
async for value in aio.as_generated(gens):
|
||||
results.append(value)
|
||||
self.assertListEqual(sorted(expected), sorted(results))
|
||||
|
||||
@async_test
|
||||
async def test_as_generated_cancelled(self):
|
||||
async def gen():
|
||||
for i in range(5):
|
||||
yield i
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
expected = [0, 0, 1, 1]
|
||||
results = []
|
||||
|
||||
async def foo():
|
||||
gens = [gen(), gen()]
|
||||
async for value in aio.as_generated(gens):
|
||||
results.append(value)
|
||||
return results
|
||||
|
||||
task = asyncio.ensure_future(foo())
|
||||
await asyncio.sleep(0.15)
|
||||
task.cancel()
|
||||
await task
|
||||
|
||||
self.assertListEqual(sorted(expected), sorted(results))
|
||||
|
||||
@async_test
|
||||
async def test_gather_input_types(self):
|
||||
async def fn(arg):
|
||||
await asyncio.sleep(0.001)
|
||||
return arg
|
||||
|
||||
fns = [fn(1), asyncio.ensure_future(fn(2))]
|
||||
if hasattr(asyncio, "create_task"):
|
||||
# 3.7 only
|
||||
fns.append(asyncio.create_task(fn(3)))
|
||||
else:
|
||||
fns.append(fn(3))
|
||||
|
||||
result = await aio.gather(*fns)
|
||||
self.assertEqual([1, 2, 3], result)
|
||||
|
||||
@async_test
|
||||
async def test_gather_limited(self):
|
||||
max_counter = 0
|
||||
counter = 0
|
||||
|
||||
async def fn(arg):
|
||||
nonlocal counter, max_counter
|
||||
counter += 1
|
||||
max_counter = max(max_counter, counter)
|
||||
await asyncio.sleep(0.001)
|
||||
counter -= 1
|
||||
return arg
|
||||
|
||||
# Limit of 2
|
||||
result = await aio.gather(*[fn(i) for i in range(10)], limit=2)
|
||||
self.assertEqual(2, max_counter)
|
||||
self.assertEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], result)
|
||||
|
||||
# No limit
|
||||
result = await aio.gather(*[fn(i) for i in range(10)])
|
||||
self.assertEqual(
|
||||
10, max_counter
|
||||
) # TODO: on a loaded machine this might be less?
|
||||
self.assertEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], result)
|
||||
|
||||
@async_test
|
||||
async def test_gather_limited_dupes(self):
|
||||
async def fn(arg):
|
||||
await asyncio.sleep(0.001)
|
||||
return arg
|
||||
|
||||
f = fn(1)
|
||||
g = fn(2)
|
||||
result = await aio.gather(f, f, f, g, f, g, limit=2)
|
||||
self.assertEqual([1, 1, 1, 2, 1, 2], result)
|
||||
|
||||
f = fn(1)
|
||||
g = fn(2)
|
||||
result = await aio.gather(f, f, f, g, f, g)
|
||||
self.assertEqual([1, 1, 1, 2, 1, 2], result)
|
||||
|
||||
@async_test
|
||||
async def test_gather_with_exceptions(self):
|
||||
class MyException(Exception):
|
||||
pass
|
||||
|
||||
async def fn(arg, fail=False):
|
||||
await asyncio.sleep(arg)
|
||||
if fail:
|
||||
raise MyException(arg)
|
||||
return arg
|
||||
|
||||
with self.assertRaises(MyException):
|
||||
await aio.gather(fn(0.002, fail=True), fn(0.001))
|
||||
|
||||
result = await aio.gather(
|
||||
fn(0.002, fail=True), fn(0.001), return_exceptions=True
|
||||
)
|
||||
self.assertEqual(result[1], 0.001)
|
||||
self.assertIsInstance(result[0], MyException)
|
||||
|
||||
@async_test
|
||||
async def test_gather_cancel(self):
|
||||
cancelled = False
|
||||
started = False
|
||||
|
||||
async def _fn():
|
||||
nonlocal started, cancelled
|
||||
try:
|
||||
started = True
|
||||
await asyncio.sleep(10) # might as well be forever
|
||||
except asyncio.CancelledError:
|
||||
nonlocal cancelled
|
||||
cancelled = True
|
||||
raise
|
||||
|
||||
async def _gather():
|
||||
await aio.gather(_fn())
|
||||
|
||||
if hasattr(asyncio, "create_task"):
|
||||
# 3.7+ only
|
||||
task = asyncio.create_task(_gather())
|
||||
else:
|
||||
task = asyncio.ensure_future(_gather())
|
||||
# to insure the gather actually runs
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await task
|
||||
self.assertTrue(started)
|
||||
self.assertTrue(cancelled)
|
||||
@@ -0,0 +1,372 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from unittest import TestCase
|
||||
|
||||
import aioitertools as ait
|
||||
from .helpers import async_test
|
||||
|
||||
slist = ["A", "B", "C"]
|
||||
srange = range(3)
|
||||
srange1 = range(1, 4)
|
||||
srange0 = range(1)
|
||||
|
||||
|
||||
class BuiltinsTest(TestCase):
|
||||
|
||||
# aioitertools.all()
|
||||
|
||||
@async_test
|
||||
async def test_all_list(self):
|
||||
self.assertTrue(await ait.all([True, 1, "string"]))
|
||||
self.assertFalse(await ait.all([True, 0, "string"]))
|
||||
|
||||
@async_test
|
||||
async def test_all_range(self):
|
||||
self.assertTrue(await ait.all(srange1))
|
||||
self.assertFalse(await ait.all(srange))
|
||||
|
||||
@async_test
|
||||
async def test_all_generator(self):
|
||||
self.assertTrue(await ait.all(x for x in srange1))
|
||||
self.assertFalse(await ait.all(x for x in srange))
|
||||
|
||||
@async_test
|
||||
async def test_all_async_generator(self):
|
||||
self.assertTrue(await ait.all(ait.iter(srange1)))
|
||||
self.assertFalse(await ait.all(ait.iter(srange)))
|
||||
|
||||
# aioitertools.any()
|
||||
|
||||
@async_test
|
||||
async def test_any_list(self):
|
||||
self.assertTrue(await ait.any([False, 1, ""]))
|
||||
self.assertFalse(await ait.any([False, 0, ""]))
|
||||
|
||||
@async_test
|
||||
async def test_any_range(self):
|
||||
self.assertTrue(await ait.any(srange))
|
||||
self.assertTrue(await ait.any(srange1))
|
||||
self.assertFalse(await ait.any(srange0))
|
||||
|
||||
@async_test
|
||||
async def test_any_generator(self):
|
||||
self.assertTrue(await ait.any(x for x in srange))
|
||||
self.assertTrue(await ait.any(x for x in srange1))
|
||||
self.assertFalse(await ait.any(x for x in srange0))
|
||||
|
||||
@async_test
|
||||
async def test_any_async_generator(self):
|
||||
self.assertTrue(await ait.any(ait.iter(srange)))
|
||||
self.assertTrue(await ait.any(ait.iter(srange1)))
|
||||
self.assertFalse(await ait.any(ait.iter(srange0)))
|
||||
|
||||
# aioitertools.iter()
|
||||
|
||||
@async_test
|
||||
async def test_iter_list(self):
|
||||
it = ait.iter(slist)
|
||||
self.assertIsInstance(it, AsyncIterator)
|
||||
idx = 0
|
||||
async for item in it:
|
||||
self.assertEqual(item, slist[idx])
|
||||
idx += 1
|
||||
|
||||
@async_test
|
||||
async def test_iter_range(self):
|
||||
it = ait.iter(srange)
|
||||
self.assertIsInstance(it, AsyncIterator)
|
||||
idx = 0
|
||||
async for item in it:
|
||||
self.assertEqual(item, srange[idx])
|
||||
idx += 1
|
||||
|
||||
@async_test
|
||||
async def test_iter_iterable(self):
|
||||
sentinel = object()
|
||||
|
||||
class async_iterable:
|
||||
def __aiter__(self):
|
||||
return sentinel
|
||||
|
||||
aiter = async_iterable()
|
||||
self.assertEqual(ait.iter(aiter), sentinel)
|
||||
|
||||
@async_test
|
||||
async def test_iter_iterator(self):
|
||||
sentinel = object()
|
||||
|
||||
class async_iterator:
|
||||
def __aiter__(self):
|
||||
return sentinel
|
||||
|
||||
def __anext__(self):
|
||||
return sentinel
|
||||
|
||||
aiter = async_iterator()
|
||||
self.assertEqual(ait.iter(aiter), aiter)
|
||||
|
||||
@async_test
|
||||
async def test_iter_async_generator(self):
|
||||
async def async_gen():
|
||||
yield 1
|
||||
yield 2
|
||||
|
||||
agen = async_gen()
|
||||
self.assertEqual(ait.iter(agen), agen)
|
||||
|
||||
# aioitertools.next()
|
||||
|
||||
@async_test
|
||||
async def test_next_list(self):
|
||||
it = ait.iter(slist)
|
||||
self.assertEqual(await ait.next(it), "A")
|
||||
self.assertEqual(await ait.next(it), "B")
|
||||
self.assertEqual(await ait.next(it), "C")
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_next_range(self):
|
||||
it = ait.iter(srange)
|
||||
self.assertEqual(await ait.next(it), 0)
|
||||
self.assertEqual(await ait.next(it), 1)
|
||||
self.assertEqual(await ait.next(it), 2)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_next_iterable(self):
|
||||
class async_iter:
|
||||
def __init__(self):
|
||||
self.index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
def __anext__(self):
|
||||
if self.index > 2:
|
||||
raise StopAsyncIteration()
|
||||
return self.fake_next()
|
||||
|
||||
async def fake_next(self):
|
||||
value = slist[self.index]
|
||||
self.index += 1
|
||||
return value
|
||||
|
||||
it = ait.iter(async_iter())
|
||||
self.assertEqual(await ait.next(it), "A")
|
||||
self.assertEqual(await ait.next(it), "B")
|
||||
self.assertEqual(await ait.next(it), "C")
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
it = iter(slist)
|
||||
self.assertEqual(await ait.next(it), "A")
|
||||
self.assertEqual(await ait.next(it), "B")
|
||||
self.assertEqual(await ait.next(it), "C")
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_next_async_generator(self):
|
||||
async def async_gen():
|
||||
for item in slist:
|
||||
yield item
|
||||
|
||||
it = ait.iter(async_gen())
|
||||
self.assertEqual(await ait.next(it), "A")
|
||||
self.assertEqual(await ait.next(it), "B")
|
||||
self.assertEqual(await ait.next(it), "C")
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_next_default_iterable(self):
|
||||
it = iter(["A"])
|
||||
|
||||
self.assertEqual(await ait.next(it, "?"), "A")
|
||||
# End of iteration
|
||||
self.assertEqual(await ait.next(it, "?"), "?")
|
||||
|
||||
@async_test
|
||||
async def test_next_default_async_iterable(self):
|
||||
it = ait.iter(["A"])
|
||||
self.assertEqual(await ait.next(it, "?"), "A")
|
||||
# End of iteration
|
||||
self.assertEqual(await ait.next(it, "?"), "?")
|
||||
|
||||
# aioitertools.list()
|
||||
|
||||
@async_test
|
||||
async def test_list(self):
|
||||
self.assertEqual(await ait.list(ait.iter(slist)), slist)
|
||||
|
||||
@async_test
|
||||
async def test_tuple(self):
|
||||
self.assertEqual(await ait.tuple(ait.iter(slist)), tuple(slist))
|
||||
|
||||
# aioitertools.set()
|
||||
|
||||
@async_test
|
||||
async def test_set(self):
|
||||
self.assertEqual(await ait.set(ait.iter(slist)), set(slist))
|
||||
|
||||
# aioitertools.enumerate()
|
||||
|
||||
@async_test
|
||||
async def test_enumerate(self):
|
||||
async for index, value in ait.enumerate(slist):
|
||||
self.assertEqual(value, slist[index])
|
||||
|
||||
@async_test
|
||||
async def test_enumerate_start(self):
|
||||
async for index, value in ait.enumerate(slist, 4):
|
||||
self.assertEqual(value, slist[index - 4])
|
||||
|
||||
# aioitertools.map()
|
||||
|
||||
@async_test
|
||||
async def test_map_function_list(self):
|
||||
idx = 0
|
||||
async for value in ait.map(str.lower, slist):
|
||||
self.assertEqual(value, slist[idx].lower())
|
||||
idx += 1
|
||||
|
||||
@async_test
|
||||
async def test_map_function_async_generator(self):
|
||||
async def gen():
|
||||
for item in slist:
|
||||
yield item
|
||||
|
||||
idx = 0
|
||||
async for value in ait.map(str.lower, gen()):
|
||||
self.assertEqual(value, slist[idx].lower())
|
||||
idx += 1
|
||||
|
||||
@async_test
|
||||
async def test_map_coroutine_list(self):
|
||||
async def double(x):
|
||||
await asyncio.sleep(0.0001)
|
||||
return x * 2
|
||||
|
||||
idx = 0
|
||||
async for value in ait.map(double, slist):
|
||||
self.assertEqual(value, slist[idx] * 2)
|
||||
idx += 1
|
||||
|
||||
@async_test
|
||||
async def test_map_coroutine_generator(self):
|
||||
async def gen():
|
||||
for item in slist:
|
||||
yield item
|
||||
|
||||
async def double(x):
|
||||
await asyncio.sleep(0.0001)
|
||||
return x * 2
|
||||
|
||||
idx = 0
|
||||
async for value in ait.map(double, gen()):
|
||||
self.assertEqual(value, slist[idx] * 2)
|
||||
idx += 1
|
||||
|
||||
# aioitertools.max()
|
||||
|
||||
@async_test
|
||||
async def test_max_basic(self):
|
||||
async def gen():
|
||||
for item in slist:
|
||||
yield item
|
||||
|
||||
self.assertEqual(await ait.max(gen()), "C")
|
||||
self.assertEqual(await ait.max(range(4)), 3)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "iterable is empty"):
|
||||
await ait.max([])
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "kwarg .+ not supported"):
|
||||
await ait.max(None, foo="foo")
|
||||
|
||||
@async_test
|
||||
async def test_max_default(self):
|
||||
self.assertEqual(await ait.max(range(2), default="x"), 1)
|
||||
self.assertEqual(await ait.max([], default="x"), "x")
|
||||
self.assertEqual(await ait.max([], default=None), None)
|
||||
|
||||
@async_test
|
||||
async def test_max_key(self):
|
||||
words = ["star", "buzz", "guard"]
|
||||
|
||||
def reverse(s):
|
||||
return s[::-1]
|
||||
|
||||
self.assertEqual(reverse("python"), "nohtyp")
|
||||
|
||||
self.assertEqual(await ait.max(words), "star")
|
||||
self.assertEqual(await ait.max(words, key=reverse), "buzz")
|
||||
|
||||
# aioitertools.min()
|
||||
|
||||
@async_test
|
||||
async def test_min_basic(self):
|
||||
async def gen():
|
||||
for item in slist:
|
||||
yield item
|
||||
|
||||
self.assertEqual(await ait.min(gen()), "A")
|
||||
self.assertEqual(await ait.min(range(4)), 0)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "iterable is empty"):
|
||||
await ait.min([])
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "kwarg .+ not supported"):
|
||||
await ait.min(None, foo="foo")
|
||||
|
||||
@async_test
|
||||
async def test_min_default(self):
|
||||
self.assertEqual(await ait.min(range(2), default="x"), 0)
|
||||
self.assertEqual(await ait.min([], default="x"), "x")
|
||||
self.assertEqual(await ait.min([], default=None), None)
|
||||
|
||||
@async_test
|
||||
async def test_min_key(self):
|
||||
words = ["star", "buzz", "guard"]
|
||||
|
||||
def reverse(s):
|
||||
return s[::-1]
|
||||
|
||||
self.assertEqual(reverse("python"), "nohtyp")
|
||||
|
||||
self.assertEqual(await ait.min(words), "buzz")
|
||||
self.assertEqual(await ait.min(words, key=reverse), "guard")
|
||||
|
||||
# aioitertools.sum()
|
||||
|
||||
@async_test
|
||||
async def test_sum_range_default(self):
|
||||
self.assertEqual(await ait.sum(srange), sum(srange))
|
||||
|
||||
@async_test
|
||||
async def test_sum_list_string(self):
|
||||
self.assertEqual(await ait.sum(slist, "foo"), "fooABC")
|
||||
|
||||
# aioitertools.zip()
|
||||
|
||||
@async_test
|
||||
async def test_zip_equal(self):
|
||||
idx = 0
|
||||
async for a, b in ait.zip(slist, srange):
|
||||
self.assertEqual(a, slist[idx])
|
||||
self.assertEqual(b, srange[idx])
|
||||
idx += 1
|
||||
|
||||
@async_test
|
||||
async def test_zip_shortest(self):
|
||||
short = ["a", "b", "c"]
|
||||
long = [0, 1, 2, 3, 5]
|
||||
|
||||
result = await ait.list(ait.zip(short, long))
|
||||
expected = [("a", 0), ("b", 1), ("c", 2)]
|
||||
self.assertListEqual(expected, result)
|
||||
@@ -0,0 +1,57 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import sys
|
||||
from unittest import skipIf, TestCase
|
||||
|
||||
from aioitertools.helpers import maybe_await
|
||||
|
||||
|
||||
def async_test(fn):
|
||||
def wrapped(*args, **kwargs):
|
||||
try:
|
||||
loop = asyncio.new_event_loop()
|
||||
loop.set_debug(False)
|
||||
result = loop.run_until_complete(fn(*args, **kwargs))
|
||||
return result
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
class HelpersTest(TestCase):
|
||||
|
||||
# aioitertools.helpers.maybe_await()
|
||||
|
||||
@async_test
|
||||
async def test_maybe_await(self):
|
||||
self.assertEqual(await maybe_await(42), 42)
|
||||
|
||||
@async_test
|
||||
async def test_maybe_await_async_def(self):
|
||||
async def forty_two():
|
||||
await asyncio.sleep(0.0001)
|
||||
return 42
|
||||
|
||||
self.assertEqual(await maybe_await(forty_two()), 42)
|
||||
|
||||
@skipIf(sys.version_info >= (3, 11), "@asyncio.coroutine removed")
|
||||
@async_test
|
||||
async def test_maybe_await_coroutine(self):
|
||||
@asyncio.coroutine
|
||||
def forty_two():
|
||||
yield from asyncio.sleep(0.0001)
|
||||
return 42
|
||||
|
||||
self.assertEqual(await maybe_await(forty_two()), 42)
|
||||
|
||||
@async_test
|
||||
async def test_maybe_await_partial(self):
|
||||
async def multiply(a, b):
|
||||
await asyncio.sleep(0.0001)
|
||||
return a * b
|
||||
|
||||
self.assertEqual(await maybe_await(functools.partial(multiply, 6)(7)), 42)
|
||||
@@ -0,0 +1,790 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
import asyncio
|
||||
import operator
|
||||
from unittest import TestCase
|
||||
|
||||
import aioitertools as ait
|
||||
from .helpers import async_test
|
||||
|
||||
slist = ["A", "B", "C"]
|
||||
srange = range(1, 4)
|
||||
|
||||
|
||||
class ItertoolsTest(TestCase):
|
||||
@async_test
|
||||
async def test_accumulate_range_default(self):
|
||||
it = ait.accumulate(srange)
|
||||
for k in [1, 3, 6]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_accumulate_range_function(self):
|
||||
it = ait.accumulate(srange, func=operator.mul)
|
||||
for k in [1, 2, 6]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_accumulate_range_coroutine(self):
|
||||
async def mul(a, b):
|
||||
return a * b
|
||||
|
||||
it = ait.accumulate(srange, func=mul)
|
||||
for k in [1, 2, 6]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_accumulate_gen_function(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 4
|
||||
|
||||
it = ait.accumulate(gen(), func=operator.mul)
|
||||
for k in [1, 2, 8]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_accumulate_gen_coroutine(self):
|
||||
async def mul(a, b):
|
||||
return a * b
|
||||
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 4
|
||||
|
||||
it = ait.accumulate(gen(), func=mul)
|
||||
for k in [1, 2, 8]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_accumulate_empty(self):
|
||||
values = []
|
||||
async for value in ait.accumulate([]):
|
||||
values.append(value)
|
||||
|
||||
self.assertEqual(values, [])
|
||||
|
||||
@async_test
|
||||
async def test_batched(self):
|
||||
test_matrix = [
|
||||
([], 1, []),
|
||||
([1, 2, 3], 1, [(1,), (2,), (3,)]),
|
||||
([2, 3, 4], 2, [(2, 3), (4,)]),
|
||||
([5, 6], 3, [(5, 6)]),
|
||||
(ait.iter([-2, -1, 0, 1, 2]), 2, [(-2, -1), (0, 1), (2,)]),
|
||||
]
|
||||
for iterable, batch_size, answer in test_matrix:
|
||||
result = [batch async for batch in ait.batched(iterable, batch_size)]
|
||||
|
||||
self.assertEqual(result, answer)
|
||||
|
||||
@async_test
|
||||
async def test_batched_errors(self):
|
||||
with self.assertRaisesRegex(ValueError, "n must be at least one"):
|
||||
[batch async for batch in ait.batched([1], 0)]
|
||||
with self.assertRaisesRegex(ValueError, "incomplete batch"):
|
||||
[batch async for batch in ait.batched([1, 2, 3], 2, strict=True)]
|
||||
|
||||
@async_test
|
||||
async def test_chain_lists(self):
|
||||
it = ait.chain(slist, srange)
|
||||
for k in ["A", "B", "C", 1, 2, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_chain_list_gens(self):
|
||||
async def gen():
|
||||
for k in range(2, 9, 2):
|
||||
yield k
|
||||
|
||||
it = ait.chain(slist, gen())
|
||||
for k in ["A", "B", "C", 2, 4, 6, 8]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_chain_from_iterable(self):
|
||||
async def gen():
|
||||
for k in range(2, 9, 2):
|
||||
yield k
|
||||
|
||||
it = ait.chain.from_iterable([slist, gen()])
|
||||
for k in ["A", "B", "C", 2, 4, 6, 8]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_chain_from_iterable_parameter_expansion_gen(self):
|
||||
async def gen():
|
||||
for k in range(2, 9, 2):
|
||||
yield k
|
||||
|
||||
async def parameters_gen():
|
||||
yield slist
|
||||
yield gen()
|
||||
|
||||
it = ait.chain.from_iterable(parameters_gen())
|
||||
for k in ["A", "B", "C", 2, 4, 6, 8]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_combinations(self):
|
||||
it = ait.combinations(range(4), 3)
|
||||
for k in [(0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3)]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_combinations_with_replacement(self):
|
||||
it = ait.combinations_with_replacement(slist, 2)
|
||||
for k in [
|
||||
("A", "A"),
|
||||
("A", "B"),
|
||||
("A", "C"),
|
||||
("B", "B"),
|
||||
("B", "C"),
|
||||
("C", "C"),
|
||||
]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_compress_list(self):
|
||||
data = range(10)
|
||||
selectors = [0, 1, 1, 0, 0, 0, 1, 0, 1, 0]
|
||||
|
||||
it = ait.compress(data, selectors)
|
||||
for k in [1, 2, 6, 8]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_compress_gen(self):
|
||||
data = "abcdefghijkl"
|
||||
selectors = ait.cycle([1, 0, 0])
|
||||
|
||||
it = ait.compress(data, selectors)
|
||||
for k in ["a", "d", "g", "j"]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_count_bare(self):
|
||||
it = ait.count()
|
||||
for k in [0, 1, 2, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
|
||||
@async_test
|
||||
async def test_count_start(self):
|
||||
it = ait.count(42)
|
||||
for k in [42, 43, 44, 45]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
|
||||
@async_test
|
||||
async def test_count_start_step(self):
|
||||
it = ait.count(42, 3)
|
||||
for k in [42, 45, 48, 51]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
|
||||
@async_test
|
||||
async def test_count_negative(self):
|
||||
it = ait.count(step=-2)
|
||||
for k in [0, -2, -4, -6]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
|
||||
@async_test
|
||||
async def test_cycle_list(self):
|
||||
it = ait.cycle(slist)
|
||||
for k in ["A", "B", "C", "A", "B", "C", "A", "B"]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
|
||||
@async_test
|
||||
async def test_cycle_gen(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 42
|
||||
|
||||
it = ait.cycle(gen())
|
||||
for k in [1, 2, 42, 1, 2, 42, 1, 2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
|
||||
@async_test
|
||||
async def test_dropwhile_empty(self):
|
||||
def pred(x):
|
||||
return x < 2
|
||||
|
||||
result = await ait.list(ait.dropwhile(pred, []))
|
||||
self.assertEqual(result, [])
|
||||
|
||||
@async_test
|
||||
async def test_dropwhile_function_list(self):
|
||||
def pred(x):
|
||||
return x < 2
|
||||
|
||||
it = ait.dropwhile(pred, srange)
|
||||
for k in [2, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_dropwhile_function_gen(self):
|
||||
def pred(x):
|
||||
return x < 2
|
||||
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 42
|
||||
|
||||
it = ait.dropwhile(pred, gen())
|
||||
for k in [2, 42]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_dropwhile_coroutine_list(self):
|
||||
async def pred(x):
|
||||
return x < 2
|
||||
|
||||
it = ait.dropwhile(pred, srange)
|
||||
for k in [2, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_dropwhile_coroutine_gen(self):
|
||||
async def pred(x):
|
||||
return x < 2
|
||||
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 42
|
||||
|
||||
it = ait.dropwhile(pred, gen())
|
||||
for k in [2, 42]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_filterfalse_function_list(self):
|
||||
def pred(x):
|
||||
return x % 2 == 0
|
||||
|
||||
it = ait.filterfalse(pred, srange)
|
||||
for k in [1, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_filterfalse_coroutine_list(self):
|
||||
async def pred(x):
|
||||
return x % 2 == 0
|
||||
|
||||
it = ait.filterfalse(pred, srange)
|
||||
for k in [1, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_groupby_list(self):
|
||||
data = "aaabba"
|
||||
|
||||
it = ait.groupby(data)
|
||||
for k in [("a", ["a", "a", "a"]), ("b", ["b", "b"]), ("a", ["a"])]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_groupby_list_key(self):
|
||||
data = "aAabBA"
|
||||
|
||||
it = ait.groupby(data, key=str.lower)
|
||||
for k in [("a", ["a", "A", "a"]), ("b", ["b", "B"]), ("a", ["A"])]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_groupby_gen(self):
|
||||
async def gen():
|
||||
for c in "aaabba":
|
||||
yield c
|
||||
|
||||
it = ait.groupby(gen())
|
||||
for k in [("a", ["a", "a", "a"]), ("b", ["b", "b"]), ("a", ["a"])]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_groupby_gen_key(self):
|
||||
async def gen():
|
||||
for c in "aAabBA":
|
||||
yield c
|
||||
|
||||
it = ait.groupby(gen(), key=str.lower)
|
||||
for k in [("a", ["a", "A", "a"]), ("b", ["b", "B"]), ("a", ["A"])]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_groupby_empty(self):
|
||||
async def gen():
|
||||
for _ in range(0):
|
||||
yield # Force generator with no actual iteration
|
||||
|
||||
async for _ in ait.groupby(gen()):
|
||||
self.fail("No iteration should have happened")
|
||||
|
||||
@async_test
|
||||
async def test_islice_bad_range(self):
|
||||
with self.assertRaisesRegex(ValueError, "must pass stop index"):
|
||||
async for _ in ait.islice([1, 2]):
|
||||
pass
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "too many arguments"):
|
||||
async for _ in ait.islice([1, 2], 1, 2, 3, 4):
|
||||
pass
|
||||
|
||||
@async_test
|
||||
async def test_islice_stop_zero(self):
|
||||
values = []
|
||||
async for value in ait.islice(range(5), 0):
|
||||
values.append(value)
|
||||
self.assertEqual(values, [])
|
||||
|
||||
@async_test
|
||||
async def test_islice_range_stop(self):
|
||||
it = ait.islice(srange, 2)
|
||||
for k in [1, 2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_islice_range_start_step(self):
|
||||
it = ait.islice(srange, 0, None, 2)
|
||||
for k in [1, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_islice_range_start_stop(self):
|
||||
it = ait.islice(srange, 1, 3)
|
||||
for k in [2, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_islice_range_start_stop_step(self):
|
||||
it = ait.islice(srange, 1, 3, 2)
|
||||
for k in [2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_islice_gen_stop(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 3
|
||||
yield 4
|
||||
|
||||
gen_it = gen()
|
||||
it = ait.islice(gen_it, 2)
|
||||
for k in [1, 2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
assert await ait.list(gen_it) == [3, 4]
|
||||
|
||||
@async_test
|
||||
async def test_islice_gen_start_step(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 3
|
||||
yield 4
|
||||
|
||||
it = ait.islice(gen(), 1, None, 2)
|
||||
for k in [2, 4]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_islice_gen_start_stop(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 3
|
||||
yield 4
|
||||
|
||||
it = ait.islice(gen(), 1, 3)
|
||||
for k in [2, 3]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_islice_gen_start_stop_step(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 3
|
||||
yield 4
|
||||
|
||||
gen_it = gen()
|
||||
it = ait.islice(gen_it, 1, 3, 2)
|
||||
for k in [2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
assert await ait.list(gen_it) == [4]
|
||||
|
||||
@async_test
|
||||
async def test_permutations_list(self):
|
||||
it = ait.permutations(srange, r=2)
|
||||
for k in [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_permutations_gen(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 3
|
||||
|
||||
it = ait.permutations(gen(), r=2)
|
||||
for k in [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_product_list(self):
|
||||
it = ait.product([1, 2], [6, 7])
|
||||
for k in [(1, 6), (1, 7), (2, 6), (2, 7)]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_product_gen(self):
|
||||
async def gen(x):
|
||||
yield x
|
||||
yield x + 1
|
||||
|
||||
it = ait.product(gen(1), gen(6))
|
||||
for k in [(1, 6), (1, 7), (2, 6), (2, 7)]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_repeat(self):
|
||||
it = ait.repeat(42)
|
||||
for k in [42] * 10:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
|
||||
@async_test
|
||||
async def test_repeat_limit(self):
|
||||
it = ait.repeat(42, 5)
|
||||
for k in [42] * 5:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_starmap_function_list(self):
|
||||
data = [slist[:2], slist[1:], slist]
|
||||
|
||||
def concat(*args):
|
||||
return "".join(args)
|
||||
|
||||
it = ait.starmap(concat, data)
|
||||
for k in ["AB", "BC", "ABC"]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_starmap_function_gen(self):
|
||||
def gen():
|
||||
yield slist[:2]
|
||||
yield slist[1:]
|
||||
yield slist
|
||||
|
||||
def concat(*args):
|
||||
return "".join(args)
|
||||
|
||||
it = ait.starmap(concat, gen())
|
||||
for k in ["AB", "BC", "ABC"]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_starmap_coroutine_list(self):
|
||||
data = [slist[:2], slist[1:], slist]
|
||||
|
||||
async def concat(*args):
|
||||
return "".join(args)
|
||||
|
||||
it = ait.starmap(concat, data)
|
||||
for k in ["AB", "BC", "ABC"]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_starmap_coroutine_gen(self):
|
||||
async def gen():
|
||||
yield slist[:2]
|
||||
yield slist[1:]
|
||||
yield slist
|
||||
|
||||
async def concat(*args):
|
||||
return "".join(args)
|
||||
|
||||
it = ait.starmap(concat, gen())
|
||||
for k in ["AB", "BC", "ABC"]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_takewhile_empty(self):
|
||||
def pred(x):
|
||||
return x < 3
|
||||
|
||||
values = await ait.list(ait.takewhile(pred, []))
|
||||
self.assertEqual(values, [])
|
||||
|
||||
@async_test
|
||||
async def test_takewhile_function_list(self):
|
||||
def pred(x):
|
||||
return x < 3
|
||||
|
||||
it = ait.takewhile(pred, srange)
|
||||
for k in [1, 2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_takewhile_function_gen(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 3
|
||||
|
||||
def pred(x):
|
||||
return x < 3
|
||||
|
||||
it = ait.takewhile(pred, gen())
|
||||
for k in [1, 2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_takewhile_coroutine_list(self):
|
||||
async def pred(x):
|
||||
return x < 3
|
||||
|
||||
it = ait.takewhile(pred, srange)
|
||||
for k in [1, 2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_takewhile_coroutine_gen(self):
|
||||
def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
yield 3
|
||||
|
||||
async def pred(x):
|
||||
return x < 3
|
||||
|
||||
it = ait.takewhile(pred, gen())
|
||||
for k in [1, 2]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_tee_list_two(self):
|
||||
it1, it2 = ait.tee(slist * 2)
|
||||
|
||||
for k in slist * 2:
|
||||
a, b = await asyncio.gather(ait.next(it1), ait.next(it2))
|
||||
self.assertEqual(a, b)
|
||||
self.assertEqual(a, k)
|
||||
self.assertEqual(b, k)
|
||||
for it in [it1, it2]:
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_tee_list_six(self):
|
||||
itrs = ait.tee(slist * 2, n=6)
|
||||
|
||||
for k in slist * 2:
|
||||
values = await asyncio.gather(*[ait.next(it) for it in itrs])
|
||||
for value in values:
|
||||
self.assertEqual(value, k)
|
||||
for it in itrs:
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_tee_gen_two(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 4
|
||||
yield 9
|
||||
yield 16
|
||||
|
||||
it1, it2 = ait.tee(gen())
|
||||
|
||||
for k in [1, 4, 9, 16]:
|
||||
a, b = await asyncio.gather(ait.next(it1), ait.next(it2))
|
||||
self.assertEqual(a, b)
|
||||
self.assertEqual(a, k)
|
||||
self.assertEqual(b, k)
|
||||
for it in [it1, it2]:
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_tee_gen_six(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 4
|
||||
yield 9
|
||||
yield 16
|
||||
|
||||
itrs = ait.tee(gen(), n=6)
|
||||
|
||||
for k in [1, 4, 9, 16]:
|
||||
values = await asyncio.gather(*[ait.next(it) for it in itrs])
|
||||
for value in values:
|
||||
self.assertEqual(value, k)
|
||||
for it in itrs:
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_tee_propagate_exception(self):
|
||||
class MyError(Exception):
|
||||
pass
|
||||
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
raise MyError
|
||||
|
||||
async def consumer(it):
|
||||
result = 0
|
||||
async for item in it:
|
||||
result += item
|
||||
return result
|
||||
|
||||
it1, it2 = ait.tee(gen())
|
||||
|
||||
values = await asyncio.gather(
|
||||
consumer(it1),
|
||||
consumer(it2),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
for value in values:
|
||||
self.assertIsInstance(value, MyError)
|
||||
|
||||
@async_test
|
||||
async def test_zip_longest_range(self):
|
||||
a = range(3)
|
||||
b = range(5)
|
||||
|
||||
it = ait.zip_longest(a, b)
|
||||
|
||||
for k in [(0, 0), (1, 1), (2, 2), (None, 3), (None, 4)]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_zip_longest_fillvalue(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 4
|
||||
yield 9
|
||||
yield 16
|
||||
|
||||
a = gen()
|
||||
b = range(5)
|
||||
|
||||
it = ait.zip_longest(a, b, fillvalue=42)
|
||||
|
||||
for k in [(1, 0), (4, 1), (9, 2), (16, 3), (42, 4)]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaises(StopAsyncIteration):
|
||||
await ait.next(it)
|
||||
|
||||
@async_test
|
||||
async def test_zip_longest_exception(self):
|
||||
async def gen():
|
||||
yield 1
|
||||
yield 2
|
||||
raise Exception("fake error")
|
||||
|
||||
a = gen()
|
||||
b = ait.repeat(5)
|
||||
|
||||
it = ait.zip_longest(a, b)
|
||||
|
||||
for k in [(1, 5), (2, 5)]:
|
||||
self.assertEqual(await ait.next(it), k)
|
||||
with self.assertRaisesRegex(Exception, "fake error"):
|
||||
await ait.next(it)
|
||||
@@ -0,0 +1,96 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
from collections.abc import AsyncIterable
|
||||
from unittest import TestCase
|
||||
|
||||
import aioitertools.more_itertools as mit
|
||||
from .helpers import async_test
|
||||
|
||||
|
||||
async def _gen() -> AsyncIterable[int]:
|
||||
for i in range(5):
|
||||
yield i
|
||||
|
||||
|
||||
async def _empty() -> AsyncIterable[int]:
|
||||
return
|
||||
yield 0
|
||||
|
||||
|
||||
class MoreItertoolsTest(TestCase):
|
||||
@async_test
|
||||
async def test_take(self) -> None:
|
||||
self.assertEqual(await mit.take(2, _gen()), [0, 1])
|
||||
self.assertEqual(await mit.take(2, range(5)), [0, 1])
|
||||
|
||||
@async_test
|
||||
async def test_take_zero(self) -> None:
|
||||
self.assertEqual(await mit.take(0, _gen()), [])
|
||||
|
||||
@async_test
|
||||
async def test_take_negative(self) -> None:
|
||||
with self.assertRaises(ValueError):
|
||||
await mit.take(-1, _gen())
|
||||
|
||||
@async_test
|
||||
async def test_take_more_than_iterable(self) -> None:
|
||||
self.assertEqual(await mit.take(10, _gen()), list(range(5)))
|
||||
|
||||
@async_test
|
||||
async def test_take_empty(self) -> None:
|
||||
it = _gen()
|
||||
self.assertEqual(len(await mit.take(5, it)), 5)
|
||||
self.assertEqual(await mit.take(1, it), [])
|
||||
self.assertEqual(await mit.take(1, _empty()), [])
|
||||
|
||||
@async_test
|
||||
async def test_chunked(self) -> None:
|
||||
self.assertEqual(
|
||||
[chunk async for chunk in mit.chunked(_gen(), 2)], [[0, 1], [2, 3], [4]]
|
||||
)
|
||||
self.assertEqual(
|
||||
[chunk async for chunk in mit.chunked(range(5), 2)], [[0, 1], [2, 3], [4]]
|
||||
)
|
||||
|
||||
@async_test
|
||||
async def test_chunked_empty(self) -> None:
|
||||
self.assertEqual([], [chunk async for chunk in mit.chunked(_empty(), 2)])
|
||||
|
||||
@async_test
|
||||
async def test_before_and_after_split(self) -> None:
|
||||
it = _gen()
|
||||
before, after = await mit.before_and_after(lambda i: i <= 2, it)
|
||||
self.assertEqual([elm async for elm in before], [0, 1, 2])
|
||||
self.assertEqual([elm async for elm in after], [3, 4])
|
||||
|
||||
@async_test
|
||||
async def test_before_and_after_before_only(self) -> None:
|
||||
it = _gen()
|
||||
before, after = await mit.before_and_after(lambda i: True, it)
|
||||
self.assertEqual([elm async for elm in before], [0, 1, 2, 3, 4])
|
||||
self.assertEqual([elm async for elm in after], [])
|
||||
|
||||
@async_test
|
||||
async def test_before_and_after_after_only(self) -> None:
|
||||
it = _gen()
|
||||
before, after = await mit.before_and_after(lambda i: False, it)
|
||||
self.assertEqual([elm async for elm in before], [])
|
||||
self.assertEqual([elm async for elm in after], [0, 1, 2, 3, 4])
|
||||
|
||||
@async_test
|
||||
async def test_before_and_after_async_predicate(self) -> None:
|
||||
async def predicate(elm: int) -> bool:
|
||||
return elm <= 2
|
||||
|
||||
it = _gen()
|
||||
before, after = await mit.before_and_after(predicate, it)
|
||||
self.assertEqual([elm async for elm in before], [0, 1, 2])
|
||||
self.assertEqual([elm async for elm in after], [3, 4])
|
||||
|
||||
@async_test
|
||||
async def test_before_and_after_empty(self) -> None:
|
||||
it = _empty()
|
||||
before, after = await mit.before_and_after(lambda i: True, it)
|
||||
self.assertEqual([elm async for elm in before], [])
|
||||
self.assertEqual([elm async for elm in after], [])
|
||||
@@ -0,0 +1,33 @@
|
||||
# Copyright 2022 Amethyst Reese
|
||||
# Licensed under the MIT license
|
||||
|
||||
import sys
|
||||
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterable, Iterator
|
||||
|
||||
from typing import Callable, TypeVar, Union
|
||||
|
||||
if sys.version_info < (3, 10): # pragma: no cover
|
||||
from typing_extensions import ParamSpec
|
||||
else: # pragma: no cover
|
||||
from typing import ParamSpec
|
||||
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
T = TypeVar("T")
|
||||
T1 = TypeVar("T1")
|
||||
T2 = TypeVar("T2")
|
||||
T3 = TypeVar("T3")
|
||||
T4 = TypeVar("T4")
|
||||
T5 = TypeVar("T5")
|
||||
N = TypeVar("N", int, float)
|
||||
|
||||
AnyFunction = Union[Callable[..., R], Callable[..., Awaitable[R]]]
|
||||
AnyIterable = Union[Iterable[T], AsyncIterable[T]]
|
||||
AnyIterableIterable = Union[Iterable[AnyIterable[T]], AsyncIterable[AnyIterable[T]]]
|
||||
AnyIterator = Union[Iterator[T], AsyncIterator[T]]
|
||||
AnyStop = (StopIteration, StopAsyncIteration)
|
||||
Accumulator = Union[Callable[[T, T], T], Callable[[T, T], Awaitable[T]]]
|
||||
KeyFunction = Union[Callable[[T], R], Callable[[T], Awaitable[R]]]
|
||||
Predicate = Union[Callable[[T], object], Callable[[T], Awaitable[object]]]
|
||||
MaybeAwaitable = Union[T, Awaitable[T]]
|
||||
Reference in New Issue
Block a user