Speedup regr_test.py by running test cases concurrently

tests/regr_test.py

@@ -4,12 +4,17 @@
 from __future__ import annotations
 
 import argparse
+import concurrent.futures
 import os
+import queue
 import re
 import shutil
 import subprocess
 import sys
 import tempfile
+import threading
+from contextlib import ExitStack, suppress
+from dataclasses import dataclass
 from enum import IntEnum
 from itertools import product
 from pathlib import Path
@@ -24,7 +29,6 @@
  get_mypy_req,
  make_venv,
  print_error,
- print_success_msg,
  testcase_dir_from_package_name,
 )
 
@@ -103,18 +107,20 @@ class Verbosity(IntEnum):
  ),
 )
 
+_PRINT_QUEUE: queue.SimpleQueue[str] = queue.SimpleQueue()
+
 
 def verbose_log(msg: str) -> None:
- print(colored("\n" + msg, "blue"))
+ _PRINT_QUEUE.put(colored(msg, "blue"))
 
 
-def setup_testcase_dir(package: PackageInfo, tempdir: Path, new_test_case_dir: Path, verbosity: Verbosity) -> None:
+def setup_testcase_dir(package: PackageInfo, tempdir: Path, verbosity: Verbosity) -> None:
  if verbosity is verbosity.VERBOSE:
- verbose_log(f"Setting up testcase dir in {tempdir}")
+ verbose_log(f"{package.name}: Setting up testcase dir in {tempdir}")
  # --warn-unused-ignores doesn't work for files inside typeshed.
  # SO, to work around this, we copy the test_cases directory into a TemporaryDirectory,
  # and run the test cases inside of that.
- shutil.copytree(package.test_case_directory, new_test_case_dir)
+ shutil.copytree(package.test_case_directory, tempdir / TEST_CASES)
  if package.is_stdlib:
  return
 
@@ -137,16 +143,15 @@ def setup_testcase_dir(package: PackageInfo, tempdir: Path, new_test_case_dir: P
  shutil.copytree(Path("stubs", requirement), new_typeshed / "stubs" / requirement)
 
  if requirements.external_pkgs:
- if verbosity is Verbosity.VERBOSE:
- verbose_log(f"Setting up venv in {tempdir / VENV_DIR}")
  pip_exe = make_venv(tempdir / VENV_DIR).pip_exe
- pip_command = [pip_exe, "install", get_mypy_req(), *requirements.external_pkgs]
+ # Use --no-cache-dir to avoid issues with concurrent read/writes to the cache
+ pip_command = [pip_exe, "install", get_mypy_req(), *requirements.external_pkgs, "--no-cache-dir"]
  if verbosity is Verbosity.VERBOSE:
- verbose_log(f"{pip_command=}")
+ verbose_log(f"{package.name}: Setting up venv in {tempdir / VENV_DIR}. {pip_command=}\n")
  try:
  subprocess.run(pip_command, check=True, capture_output=True, text=True)
  except subprocess.CalledProcessError as e:
- print(e.stderr)
+ _PRINT_QUEUE.put(f"{package.name}\n{e.stderr}")
  raise
 
 
@@ -155,10 +160,6 @@ def run_testcases(
 ) -> subprocess.CompletedProcess[str]:
  env_vars = dict(os.environ)
  new_test_case_dir = tempdir / TEST_CASES
- testcasedir_already_setup = new_test_case_dir.exists() and new_test_case_dir.is_dir()
-
- if not testcasedir_already_setup:
- setup_testcase_dir(package, tempdir=tempdir, new_test_case_dir=new_test_case_dir, verbosity=verbosity)
 
  # "--enable-error-code ignore-without-code" is purposefully omitted.
  # See https://github.com/python/typeshed/pull/8083
@@ -202,39 +203,103 @@ def run_testcases(
 
  mypy_command = [python_exe, "-m", "mypy"] + flags
  if verbosity is Verbosity.VERBOSE:
- verbose_log(f"{mypy_command=}")
+ description = f"{package.name}/{version}/{platform}"
+ msg = f"{description}: {mypy_command=}\n"
  if "MYPYPATH" in env_vars:
- verbose_log(f"{env_vars['MYPYPATH']=}")
+ msg += f"{description}: {env_vars['MYPYPATH']=}"
  else:
- verbose_log("MYPYPATH not set")
+ msg += f"{description}: MYPYPATH not set"
+ msg += "\n"
+ verbose_log(msg)
  return subprocess.run(mypy_command, capture_output=True, text=True, env=env_vars)
 
 
-def test_testcase_directory(
- package: PackageInfo, version: str, platform: str, *, verbosity: Verbosity, tempdir: Path
-) -> ReturnCode:
- msg = f"Running mypy --platform {platform} --python-version {version} on the "
- msg += "standard library test cases..." if package.is_stdlib else f"test cases for {package.name!r}..."
+@dataclass(frozen=True)
+class Result:
+ code: int
+ command_run: str
+ stderr: str
+ stdout: str
+ test_case_dir: Path
+ tempdir: Path
+
+ def print_description(self, *, verbosity: Verbosity) -> None:
+ if self.code:
+ print(f"{self.command_run}:", end=" ")
+ print_error("FAILURE\n")
+ replacements = (str(self.tempdir / TEST_CASES), str(self.test_case_dir))
+ if self.stderr:
+ print_error(self.stderr, fix_path=replacements)
+ if self.stdout:
+ print_error(self.stdout, fix_path=replacements)
+
+
+def test_testcase_directory(package: PackageInfo, version: str, platform: str, *, verbosity: Verbosity, tempdir: Path) -> Result:
+ msg = f"mypy --platform {platform} --python-version {version} on the "
+ msg += "standard library test cases" if package.is_stdlib else f"test cases for {package.name!r}"
  if verbosity > Verbosity.QUIET:
- print(msg, end=" ", flush=True)
-
- result = run_testcases(package=package, version=version, platform=platform, tempdir=tempdir, verbosity=verbosity)
-
- if result.returncode:
- if verbosity is Verbosity.QUIET:
- # We'll already have printed this if --verbosity QUIET wasn't passed.
- # If --verbosity QUIET was passed, only print this if there were errors.
- # If there are errors, the output is inscrutable if this isn't printed.
- print(msg, end=" ")
- print_error("failure\n")
- replacements = (str(tempdir / TEST_CASES), str(package.test_case_directory))
- if result.stderr:
- print_error(result.stderr, fix_path=replacements)
- if result.stdout:
- print_error(result.stdout, fix_path=replacements)
- elif verbosity > Verbosity.QUIET:
- print_success_msg()
- return result.returncode
+ _PRINT_QUEUE.put(f"Running {msg}...")
+
+ proc_info = run_testcases(package=package, version=version, platform=platform, tempdir=tempdir, verbosity=verbosity)
+ return Result(
+ code=proc_info.returncode,
+ command_run=msg,
+ stderr=proc_info.stderr,
+ stdout=proc_info.stdout,
+ test_case_dir=package.test_case_directory,
+ tempdir=tempdir,
+ )
+
+
+def print_queued_messages(ev: threading.Event) -> None:
+ while not ev.is_set():
+ with suppress(queue.Empty):
+ print(_PRINT_QUEUE.get(timeout=0.5), flush=True)
+ while True:
+ try:
+ msg = _PRINT_QUEUE.get_nowait()
+ except queue.Empty:
+ return
+ else:
+ print(msg, flush=True)
+
+
+def concurrently_run_testcases(
+ stack: ExitStack,
+ testcase_directories: list[PackageInfo],
+ verbosity: Verbosity,
+ platforms_to_test: list[str],
+ versions_to_test: list[str],
+) -> list[Result]:
+ packageinfo_to_tempdir = {
+ package_info: Path(stack.enter_context(tempfile.TemporaryDirectory())) for package_info in testcase_directories
+ }
+
+ event = threading.Event()
+ printer_thread = threading.Thread(target=print_queued_messages, args=(event,))
+ printer_thread.start()
+
+ with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) 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
+ testcase_futures = [
+ executor.submit(setup_testcase_dir, package, tempdir, verbosity)
+ for package, tempdir in packageinfo_to_tempdir.items()
+ ]
+ concurrent.futures.wait(testcase_futures)
+
+ mypy_futures = [
+ executor.submit(test_testcase_directory, testcase_dir, version, platform, verbosity=verbosity, tempdir=tempdir)
+ for (testcase_dir, tempdir), platform, version in product(
+ packageinfo_to_tempdir.items(), platforms_to_test, versions_to_test
+ )
+ ]
+ results = [future.result() for future in mypy_futures]
+
+ event.set()
+ printer_thread.join()
+ return results
 
 
 def main() -> ReturnCode:
@@ -253,16 +318,23 @@ def main() -> ReturnCode:
  versions_to_test = args.versions_to_test or [f"3.{sys.version_info[1]}"]
 
  code = 0
- for testcase_dir in testcase_directories:
- with tempfile.TemporaryDirectory() as td:
- tempdir = Path(td)
- for platform, version in product(platforms_to_test, versions_to_test):
- this_code = test_testcase_directory(testcase_dir, version, platform, verbosity=verbosity, tempdir=tempdir)
- code = max(code, this_code)
+ results: list[Result] | None = None
+
+ with ExitStack() as stack:
+ results = concurrently_run_testcases(stack, testcase_directories, verbosity, platforms_to_test, versions_to_test)
+
+ assert results is not None
+ print()
+
+ for result in results:
+ result.print_description(verbosity=verbosity)
+
+ code = max(result.code for result in results)
+
  if code:
- print_error("\nTest completed with errors")
+ print_error("Test completed with errors")
  else:
- print(colored("\nTest completed successfully!", "green"))
+ print(colored("Test completed successfully!", "green"))
 
  return code