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)]
