Allow specifying `--workers` via environment variable (#3743)

[etc/vim.git] / src / black / concurrency.py

diff --git a/src/black/concurrency.py b/src/black/concurrency.py
index bdc368d5add243b17803fa7e6bb2ff23036c6411..893eba6675a738a7b1eab2471075caa87130430a 100644 (file)
--- a/src/black/concurrency.py
+++ b/src/black/concurrency.py
@@ -47,23 +47,14 @@ def cancel(tasks: Iterable["asyncio.Task[Any]"]) -> None:
 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()]
+        to_cancel = [task for task in asyncio.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)
-            )
+        loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True))
     finally:
         # `concurrent.futures.Future` objects cannot be cancelled once they
         # are already running. There might be some when the `shutdown()` happened.
@@ -89,7 +80,8 @@ def reformat_many(
 
     executor: Executor
     if workers is None:
-        workers = os.cpu_count() or 1
+        workers = int(os.environ.get("BLACK_NUM_WORKERS", 0))
+        workers = workers or os.cpu_count() or 1
     if sys.platform == "win32":
         # Work around https://bugs.python.org/issue26903
         workers = min(workers, 60)
@@ -191,9 +183,6 @@ async def schedule_formatting(
                     sources_to_cache.append(src)
                 report.done(src, changed)
     if cancelled:
-        if sys.version_info >= (3, 7):
-            await asyncio.gather(*cancelled, return_exceptions=True)
-        else:
-            await asyncio.gather(*cancelled, loop=loop, return_exceptions=True)
+        await asyncio.gather(*cancelled, return_exceptions=True)
     if sources_to_cache:
         write_cache(cache, sources_to_cache, mode)