chore: 添加虚拟环境到仓库
- 添加 backend_service/venv 虚拟环境 - 包含所有Python依赖包 - 注意:虚拟环境约393MB,包含12655个文件
This commit is contained in:
@@ -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], [])
|
||||
Reference in New Issue
Block a user