Asset-Compression-Manager/acm/asyncio.py

import asyncio
import logging
import typing

LOG = logging.getLogger("acm.async")

async def run_command_shell(
    command,
    stdout=asyncio.subprocess.PIPE,
    stderr=asyncio.subprocess.PIPE,
    on_success: typing.List[typing.Callable] = [()]
):
    """Run command in subprocess (shell).

    Note:
        This can be used if you wish to execute e.g. "copy"
        on Windows, which can only be executed in the shell.
    """
    process = await asyncio.create_subprocess_shell(
        command, stdout=stdout, stderr=stderr
    )

    process_stdout, process_stderr = await process.communicate()

    if process.returncode == 0:
        for success_callable in on_success:
            success_callable()

    if stdout != asyncio.subprocess.DEVNULL:
        result = process_stdout.decode().strip()
        return result
    else:
        return None


def make_chunks(tasks, chunk_size):
    """Yield successive chunk_size-sized chunks from tasks.

    Note:
        Taken from https://stackoverflow.com/a/312464
        modified for python 3 only
    """
    for i in range(0, len(tasks), chunk_size):
        yield tasks[i: i + chunk_size]


def run_asyncio_commands(tasks, max_concurrent_tasks=0):
    """Run tasks asynchronously using asyncio and return results.

    If max_concurrent_tasks are set to 0, no limit is applied.

    Note:
        By default, Windows uses SelectorEventLoop, which does not support
        subprocesses. Therefore ProactorEventLoop is used on Windows.
        https://docs.python.org/3/library/asyncio-eventloops.html#windows
    """
    all_results = []

    if max_concurrent_tasks == 0:
        chunks = [tasks]
        num_chunks = len(chunks)
    else:
        chunks = make_chunks(tasks=tasks, chunk_size=max_concurrent_tasks)
        num_chunks = len(
            list(make_chunks(tasks=tasks, chunk_size=max_concurrent_tasks)))

    if asyncio.get_event_loop().is_closed():
        asyncio.set_event_loop(asyncio.new_event_loop())
    if platform.system() == "Windows":
        asyncio.set_event_loop(asyncio.ProactorEventLoop())
    loop = asyncio.get_event_loop()

    chunk = 1
    for tasks_in_chunk in chunks:
        commands = asyncio.gather(*tasks_in_chunk)
        results = loop.run_until_complete(commands)
        all_results += results
        chunk += 1

    loop.close()
    return all_results