src/black/concurrency.py

   1 import asyncio
   2 import logging
   3 import sys
   4 from typing import Any, Iterable
   5
   6 from black.output import err
   7
   8
   9 def maybe_install_uvloop() -> None:
  10     """If our environment has uvloop installed we use it.
  11
  12     This is called only from command-line entry points to avoid
  13     interfering with the parent process if Black is used as a library.
  14
  15     """
  16     try:
  17         import uvloop
  18
  19         uvloop.install()
  20     except ImportError:
  21         pass
  22
  23
  24 def cancel(tasks: Iterable["asyncio.Task[Any]"]) -> None:
  25     """asyncio signal handler that cancels all `tasks` and reports to stderr."""
  26     err("Aborted!")
  27     for task in tasks:
  28         task.cancel()
  29
  30
  31 def shutdown(loop: asyncio.AbstractEventLoop) -> None:
  32     """Cancel all pending tasks on `loop`, wait for them, and close the loop."""
  33     try:
  34         if sys.version_info[:2] >= (3, 7):
  35             all_tasks = asyncio.all_tasks
  36         else:
  37             all_tasks = asyncio.Task.all_tasks
  38         # This part is borrowed from asyncio/runners.py in Python 3.7b2.
  39         to_cancel = [task for task in all_tasks(loop) if not task.done()]
  40         if not to_cancel:
  41             return
  42
  43         for task in to_cancel:
  44             task.cancel()
  45         if sys.version_info >= (3, 7):
  46             loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True))
  47         else:
  48             loop.run_until_complete(
  49                 asyncio.gather(*to_cancel, loop=loop, return_exceptions=True)
  50             )
  51     finally:
  52         # `concurrent.futures.Future` objects cannot be cancelled once they
  53         # are already running. There might be some when the `shutdown()` happened.
  54         # Silence their logger's spew about the event loop being closed.
  55         cf_logger = logging.getLogger("concurrent.futures")
  56         cf_logger.setLevel(logging.CRITICAL)
  57         loop.close()