Speed up mypy regr test in CI

Author: JelleZijlstra

.github/workflows/tests.yml

@@ -78,7 +78,7 @@ jobs:
         with:
           version-file: "requirements-tests.txt"
       - run: uv pip install -r requirements-tests.txt --system
-      - run: python ./tests/regr_test.py --all --verbosity QUIET
+      - run: python ./tests/regr_test.py --all --verbosity QUIET --max-workers=8
 
   pyright:
     name: "pyright: Run test cases"

tests/regr_test.py

@@ -107,6 +107,13 @@ class Verbosity(IntEnum):
         "Note that this cannot be specified if --all is also specified."
     ),
 )
+parser.add_argument(
+    "-j",
+    "--max-workers",
+    type=int,
+    default=os.cpu_count(),
+    help="Maximum number of packages or mypy subprocesses to run concurrently",
+)
 
 _PRINT_QUEUE: queue.SimpleQueue[str] = queue.SimpleQueue()
 
@@ -298,6 +305,7 @@ def concurrently_run_testcases(
     verbosity: Verbosity,
     platforms_to_test: list[str],
     versions_to_test: list[str],
+    max_workers: int | None,
 ) -> list[Result]:
     packageinfo_to_tempdir = {
         distribution_info: Path(stack.enter_context(tempfile.TemporaryDirectory())) for distribution_info in testcase_directories
@@ -348,7 +356,7 @@ def cleanup_threads(
     printer_thread = threading.Thread(target=print_queued_messages, args=(event,))
     printer_thread.start()
 
-    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
         # Each temporary directory may be used by multiple processes concurrently during the next step;
         # must make sure that they're all setup correctly before starting the next step,
         # in order to avoid race conditions
@@ -358,7 +366,8 @@ def cleanup_threads(
         ]
 
         with cleanup_threads(event, printer_thread, executor):
-            concurrent.futures.wait(testcase_futures)
+            for future in concurrent.futures.as_completed(testcase_futures):
+                future.result()
 
         mypy_futures = [executor.submit(task) for task in to_do]
 
@@ -372,6 +381,8 @@ def cleanup_threads(
 
 def main() -> ReturnCode:
     args = parser.parse_args()
+    if args.max_workers is not None and args.max_workers < 1:
+        parser.error("--max-workers must be a positive integer")
 
     testcase_directories = args.packages_to_test or get_all_testcase_directories()
     verbosity = Verbosity[args.verbosity]
@@ -388,7 +399,9 @@ def main() -> ReturnCode:
     results: list[Result] | None = None
 
     with ExitStack() as stack:
-        results = concurrently_run_testcases(stack, testcase_directories, verbosity, platforms_to_test, versions_to_test)
+        results = concurrently_run_testcases(
+            stack, testcase_directories, verbosity, platforms_to_test, versions_to_test, args.max_workers
+        )
 
     assert results is not None
     if not results: