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)