async def merge(aiters):
aiters = aioitertools.iter(aiters)
iters = [None]
nexts = [asyncio.create_task(aioitertools.next(aiters))]
while len(iters) > 0:
await asyncio.wait(nexts, return_when=asyncio.FIRST_COMPLETED)
new_iters = []
completed_iters = set()
for i, future in enumerate(nexts):
if future.done():
try:
if iters[i] is None:
# new iterator
new_iters.append(aioitertools.iter(future.result()))
else:
# new item
yield future.result()
nexts[i] = asyncio.create_task(
aioitertools.next(
aiters if iters[i] is None else iters[i]))
except StopAsyncIteration:
completed_iters.add(i)
except concurrent.futures.CancelledError as e:
raise concurrent.futures.CancelledError(
'Future {!r} at index {!r} was cancelled'.format(
future, i)) from e
for i in sorted(completed_iters, reverse=True):
del iters[i]
del nexts[i]
iters += new_iters
nexts += [
asyncio.create_task(aioitertools.next(new_iter))
for new_iter in new_iters
]
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 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 def wait(aiter):
aiter = aioitertools.iter(aiter)
items = []
next_future = asyncio.create_task(aioitertools.next(aiter))
while True:
if next_future.done() or not items:
try:
items.append(asyncio.create_task(await next_future))
next_future = asyncio.create_task(aioitertools.next(aiter))
except StopAsyncIteration:
break
await asyncio.wait([items[0], next_future],
return_when=asyncio.FIRST_COMPLETED)
while items and items[0].done():
yield items.pop(0).result()
for item in items:
yield await item
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 def collate(*iterables, key=lambda x: x):
iters = [aioitertools.iter(iterable) for iterable in iterables]
heap = []
nexts = []
for i, iterable in enumerate(iters):
try:
item = await aioitertools.next(iterable)
heapq.heappush(heap, (key(item), item, i))
except StopAsyncIteration:
nexts.append(None)
else:
nexts.append(asyncio.create_task(aioitertools.next(iterable)))
while any(fut is not None for fut in nexts):
k, item, i = heapq.heappop(heap)
yield item
if nexts[i] is not None:
try:
item = await nexts[i]
heapq.heappush(heap, (key(item), item, i))
except StopAsyncIteration:
nexts[i] = None
else:
nexts[i] = asyncio.create_task(aioitertools.next(iters[i]))
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)
