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 10e288f4f93573ce6f6e599eb0da808f406b6de5..893eba6675a738a7b1eab2471075caa87130430a 100644 (file)
--- a/src/black/concurrency.py
+++ b/src/black/concurrency.py
@@ -47,12 +47,8 @@ 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
 
@@ -84,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)