async def async_map(self, fn: Callable, iterables: Iterable[Iterable[Any]], timeout: Optional[Union[int, float]] = None, chunksize: int = 1): """Aysncio compatible version of :meth:`.map`.""" if chunksize > 1: list_per_argument = [list(x) for x in iterables] n = len(list_per_argument[0]) assert all(n == len(x) for x in list_per_argument) n_chunks = (n + chunksize - 1) // chunksize iterables_chunks = [list(partition(n_chunks, x)) for x in list_per_argument] iterables_chunks = [ chunk for chunk in iterables_chunks if len(chunk) > 0] fn = chunk(fn) iterables = iterables_chunks submissions = [self.async_submit(fn, *arguments) for arguments in zip(*iterables)] futures = await asyncio.gather(*submissions) fetching_tasks = [create_task(future._async_fetch_result()) for future in futures] async def async_result_or_cancel_all(future): try: return await future.async_result(timeout=timeout) except Exception as exc: for task in fetching_tasks: task.cancel() raise exc if chunksize > 1: return (val for future in futures for val in await async_result_or_cancel_all(future)) return (await async_result_or_cancel_all(future) for future in futures)
def test_partition_toofew(): assert list(partition(6, range(3))) == [ range(0, 1), range(1, 2), range(2, 3), range(3, 3), range(3, 3), range(3, 3) ]
async def async_map(self, fn: Callable, iterables: Iterable[Iterable[Any]], timeout: Optional[Union[int, float]] = None, chunksize: int = 1): """Aysncio compatible version of :meth:`.map`.""" if not iterables: return iter([]) if chunksize > 1: list_per_argument = [list(x) for x in iterables] n = len(list_per_argument[0]) assert all(n == len(x) for x in list_per_argument) n_chunks = (n + chunksize - 1) // chunksize iterables_chunks = [ list(partition(n_chunks, x)) for x in list_per_argument ] iterables_chunks = [ chunk for chunk in iterables_chunks if len(chunk) > 0 ] fn = chunk(fn) iterables = iterables_chunks submit_tasks = [ asyncio.ensure_future(self.async_submit(fn, *arguments)) for arguments in zip(*iterables) ] try: bp_futures = [await t for t in submit_tasks] except: for t in submit_tasks: if t.done() and not t.exception(): await t.result().async_cancel() elif not t.done(): t.cancel() raise async def async_result_or_cancel_all(future): try: return await future.async_result(timeout=timeout) except: await asyncio.gather( *[bp_fut.async_cancel() for bp_fut in bp_futures], return_exceptions=True) raise if chunksize > 1: return (val for future in bp_futures for val in await async_result_or_cancel_all(future)) return (await async_result_or_cancel_all(future) for future in bp_futures)
def test_partition_even_small(): assert list(partition( 3, range(3))) == [range(0, 1), range(1, 2), range(2, 3)]
def test_partition_zero_empty(): assert list(partition(0, [])) == []
def test_partition_uneven_big(): assert list(partition(2, range(9))) == [range(0, 5), range(5, 9)]
def test_partition_even_big(): assert list(partition( 3, range(9))) == [range(0, 3), range(3, 6), range(6, 9)]