Repository: psf/black
File: src/black/concurrency.py (Python)

Code before:
import asyncio
import logging
import sys
from typing import Any, Iterable

from black.output import err


def maybe_install_uvloop() -> None:
    """If our environment has uvloop installed we use it.

    This is called only from command-line entry points to avoid
    interfering with the parent process if Black is used as a library.

    """
    try:
        import uvloop

        uvloop.install()
    except ImportError:
        pass


def cancel(tasks: Iterable["asyncio.Task[Any]"]) -> None:
    """asyncio signal handler that cancels all `tasks` and reports to stderr."""
    err("Aborted!")
    for task in tasks:
        task.cancel()


def shutdown(loop: asyncio.AbstractEventLoop) -> None:
    """Cancel all pending tasks on `loop`, wait for them, and close the loop."""
    try:
        if sys.version_info[:2] >= (3, 7):
            all_tasks = asyncio.all_tasks
        else:
            all_tasks = asyncio.Task.all_tasks
        # This part is borrowed from asyncio/runners.py in Python 3.7b2.
        to_cancel = [task for task in all_tasks(loop) if not task.done()]
        if not to_cancel:
            return

        for task in to_cancel:
            task.cancel()
        if sys.version_info >= (3, 7):
            loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True))
        else:
            loop.run_until_complete(
                asyncio.gather(*to_cancel, loop=loop, return_exceptions=True)
            )
    finally:
        # `concurrent.futures.Future` objects cannot be cancelled once they
        # are already running. There might be some when the `shutdown()` happened.
        # Silence their logger's spew about the event loop being closed.
        cf_logger = logging.getLogger("concurrent.futures")
        cf_logger.setLevel(logging.CRITICAL)
        loop.close()


Review comment (refactor):
I chose to keep the helper function instead of always attempting an import of uvloop in the module scope just as insurance if we ever accidentally break the invariant `black.concurrency` should only be imported via `black.main` if there's more than one file to format.

Code after:
        else:
            loop.run_until_complete(
                asyncio.gather(*to_cancel, loop=loop, return_exceptions=True)
            )
    finally:
        # `concurrent.futures.Future` objects cannot be cancelled once they
        # are already running. There might be some when the `shutdown()` happened.
        # Silence their logger's spew about the event loop being closed.
        cf_logger = logging.getLogger("concurrent.futures")
        cf_logger.setLevel(logging.CRITICAL)
        loop.close()


# diff-shades depends on being to monkeypatch this function to operate. I know it's
# not ideal, but this shouldn't cause any issues ... hopefully. ~ichard26
@mypyc_attr(patchable=True)
def reformat_many(
    sources: Set[Path],
    fast: bool,
    write_back: WriteBack,
    mode: Mode,
    report: Report,
    workers: Optional[int],
) -> None:
    """Reformat multiple files using a ProcessPoolExecutor."""
    maybe_install_uvloop()

    executor: Executor
    if workers is None:
        workers = os.cpu_count() or 1
    if sys.platform == "win32":
        # Work around https://bugs.python.org/issue26903
        workers = min(workers, 60)
    try:
        executor = ProcessPoolExecutor(max_workers=workers)
    except (ImportError, NotImplementedError, OSError):
        # we arrive here if the underlying system does not support multi-processing
        # like in AWS Lambda or Termux, in which case we gracefully fallback to
        # a ThreadPoolExecutor with just a single worker (more workers would not do us
        # any good due to the Global Interpreter Lock)
        executor = ThreadPoolExecutor(max_workers=1)

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(
            schedule_formatting(
                sources=sources,
                fast=fast,
                write_back=write_back,
                mode=mode,