Skip to content

Commit 4d6fda9

Browse files
authored
mypy_test.py: Allow non-types dependencies (#9408)
The approach is pretty similar to the approach I took in #9382 for regr_test.py: dynamically create virtual environments for testing stubs packages in isolation. However, since mypy_test.py tests all of typeshed's stubs in succession (and since lots of typeshed's stubs packages depend on other typeshed stubs packages), the performance issues with creating a virtual environment for testing each stubs package are even more severe than with regr_test.py. I mitigate the performance issues associated with dynamically creating virtual environments in two ways: - Dynamically creating venvs is mostly slow due to I/O bottlenecks. Creating the virtual environments can be sped up dramatically by creating them concurrently in a threadpool. The same goes for pip installing the requirements into the venvs -- though unfortunately, we have to use pip with --no-cache-dir, as the pip cache gets easily corrupted if you try to do concurrent reads and writes to the pip cache. - If types-pycocotools requires numpy>=1 and types-D3DShot also requires numpy>=1 (for example), there's no need to create 2 virtual environments. The same requirements have been specified, so they can share a virtual environment between them. This means we don't have to create nearly so many virtual environments.
1 parent 2c9816e commit 4d6fda9

File tree

2 files changed

+198
-78
lines changed

2 files changed

+198
-78
lines changed

tests/mypy_test.py

+193-76
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,18 @@
44
from __future__ import annotations
55

66
import argparse
7+
import concurrent.futures
78
import os
89
import re
10+
import subprocess
911
import sys
1012
import tempfile
11-
from contextlib import redirect_stderr, redirect_stdout
13+
import time
14+
from collections import defaultdict
1215
from dataclasses import dataclass
13-
from io import StringIO
1416
from itertools import product
1517
from pathlib import Path
18+
from threading import Lock
1619
from typing import TYPE_CHECKING, Any, NamedTuple
1720

1821
if TYPE_CHECKING:
@@ -23,17 +26,22 @@
2326
import tomli
2427
from utils import (
2528
VERSIONS_RE as VERSION_LINE_RE,
29+
PackageDependencies,
30+
VenvInfo,
2631
colored,
2732
get_gitignore_spec,
33+
get_mypy_req,
2834
get_recursive_requirements,
35+
make_venv,
2936
print_error,
3037
print_success_msg,
3138
spec_matches_path,
3239
strip_comments,
3340
)
3441

42+
# Fail early if mypy isn't installed
3543
try:
36-
from mypy.api import run as mypy_run
44+
import mypy # noqa: F401
3745
except ImportError:
3846
print_error("Cannot import mypy. Did you install it?")
3947
sys.exit(1)
@@ -108,7 +116,7 @@ class TestConfig:
108116

109117
def log(args: TestConfig, *varargs: object) -> None:
110118
if args.verbose >= 2:
111-
print(*varargs)
119+
print(colored(" ".join(map(str, varargs)), "blue"))
112120

113121

114122
def match(path: Path, args: TestConfig) -> bool:
@@ -204,7 +212,19 @@ def add_configuration(configurations: list[MypyDistConf], distribution: str) ->
204212
configurations.append(MypyDistConf(module_name, values.copy()))
205213

206214

207-
def run_mypy(args: TestConfig, configurations: list[MypyDistConf], files: list[Path], *, testing_stdlib: bool) -> ReturnCode:
215+
def run_mypy(
216+
args: TestConfig,
217+
configurations: list[MypyDistConf],
218+
files: list[Path],
219+
*,
220+
testing_stdlib: bool,
221+
non_types_dependencies: bool,
222+
venv_info: VenvInfo,
223+
mypypath: str | None = None,
224+
) -> ReturnCode:
225+
env_vars = dict(os.environ)
226+
if mypypath is not None:
227+
env_vars["MYPYPATH"] = mypypath
208228
with tempfile.NamedTemporaryFile("w+") as temp:
209229
temp.write("[mypy]\n")
210230
for dist_conf in configurations:
@@ -213,57 +233,49 @@ def run_mypy(args: TestConfig, configurations: list[MypyDistConf], files: list[P
213233
temp.write(f"{k} = {v}\n")
214234
temp.flush()
215235

216-
flags = get_mypy_flags(args, temp.name, testing_stdlib=testing_stdlib)
236+
flags = [
237+
"--python-version",
238+
args.version,
239+
"--show-traceback",
240+
"--warn-incomplete-stub",
241+
"--no-error-summary",
242+
"--platform",
243+
args.platform,
244+
"--custom-typeshed-dir",
245+
str(Path(__file__).parent.parent),
246+
"--strict",
247+
# Stub completion is checked by pyright (--allow-*-defs)
248+
"--allow-untyped-defs",
249+
"--allow-incomplete-defs",
250+
"--allow-subclassing-any", # TODO: Do we still need this now that non-types dependencies are allowed? (#5768)
251+
"--enable-error-code",
252+
"ignore-without-code",
253+
"--config-file",
254+
temp.name,
255+
]
256+
if not testing_stdlib:
257+
flags.append("--explicit-package-bases")
258+
if not non_types_dependencies:
259+
flags.append("--no-site-packages")
260+
217261
mypy_args = [*flags, *map(str, files)]
262+
mypy_command = [venv_info.python_exe, "-m", "mypy"] + mypy_args
218263
if args.verbose:
219-
print("running mypy", " ".join(mypy_args))
220-
stdout_redirect, stderr_redirect = StringIO(), StringIO()
221-
with redirect_stdout(stdout_redirect), redirect_stderr(stderr_redirect):
222-
returned_stdout, returned_stderr, exit_code = mypy_run(mypy_args)
223-
224-
if exit_code:
264+
print(colored(f"running {' '.join(mypy_command)}", "blue"))
265+
result = subprocess.run(mypy_command, capture_output=True, text=True, env=env_vars)
266+
if result.returncode:
225267
print_error("failure\n")
226-
captured_stdout = stdout_redirect.getvalue()
227-
captured_stderr = stderr_redirect.getvalue()
228-
if returned_stderr:
229-
print_error(returned_stderr)
230-
if captured_stderr:
231-
print_error(captured_stderr)
232-
if returned_stdout:
233-
print_error(returned_stdout)
234-
if captured_stdout:
235-
print_error(captured_stdout, end="")
268+
if result.stdout:
269+
print_error(result.stdout)
270+
if result.stderr:
271+
print_error(result.stderr)
272+
if non_types_dependencies and args.verbose:
273+
print("Ran with the following environment:")
274+
subprocess.run([venv_info.pip_exe, "freeze", "--all"])
275+
print()
236276
else:
237277
print_success_msg()
238-
return exit_code
239-
240-
241-
def get_mypy_flags(args: TestConfig, temp_name: str, *, testing_stdlib: bool) -> list[str]:
242-
flags = [
243-
"--python-version",
244-
args.version,
245-
"--show-traceback",
246-
"--warn-incomplete-stub",
247-
"--show-error-codes",
248-
"--no-error-summary",
249-
"--platform",
250-
args.platform,
251-
"--no-site-packages",
252-
"--custom-typeshed-dir",
253-
str(Path(__file__).parent.parent),
254-
"--strict",
255-
# Stub completion is checked by pyright (--allow-*-defs)
256-
"--allow-untyped-defs",
257-
"--allow-incomplete-defs",
258-
"--allow-subclassing-any", # Needed until we can use non-types dependencies #5768
259-
"--enable-error-code",
260-
"ignore-without-code",
261-
"--config-file",
262-
temp_name,
263-
]
264-
if not testing_stdlib:
265-
flags.append("--explicit-package-bases")
266-
return flags
278+
return result.returncode
267279

268280

269281
def add_third_party_files(
@@ -298,7 +310,9 @@ class TestResults(NamedTuple):
298310
files_checked: int
299311

300312

301-
def test_third_party_distribution(distribution: str, args: TestConfig) -> TestResults:
313+
def test_third_party_distribution(
314+
distribution: str, args: TestConfig, venv_info: VenvInfo, *, non_types_dependencies: bool
315+
) -> TestResults:
302316
"""Test the stubs of a third-party distribution.
303317
304318
Return a tuple, where the first element indicates mypy's return code
@@ -313,20 +327,24 @@ def test_third_party_distribution(distribution: str, args: TestConfig) -> TestRe
313327
if not files and args.filter:
314328
return TestResults(0, 0)
315329

316-
print(f"testing {distribution} ({len(files)} files)... ", end="")
330+
print(f"testing {distribution} ({len(files)} files)... ", end="", flush=True)
317331

318332
if not files:
319333
print_error("no files found")
320334
sys.exit(1)
321335

322-
prev_mypypath = os.getenv("MYPYPATH")
323-
os.environ["MYPYPATH"] = os.pathsep.join(str(Path("stubs", dist)) for dist in seen_dists)
324-
code = run_mypy(args, configurations, files, testing_stdlib=False)
325-
if prev_mypypath is None:
326-
del os.environ["MYPYPATH"]
327-
else:
328-
os.environ["MYPYPATH"] = prev_mypypath
329-
336+
mypypath = os.pathsep.join(str(Path("stubs", dist)) for dist in seen_dists)
337+
if args.verbose:
338+
print(colored(f"\n{mypypath=}", "blue"))
339+
code = run_mypy(
340+
args,
341+
configurations,
342+
files,
343+
venv_info=venv_info,
344+
mypypath=mypypath,
345+
testing_stdlib=False,
346+
non_types_dependencies=non_types_dependencies,
347+
)
330348
return TestResults(code, len(files))
331349

332350

@@ -343,19 +361,105 @@ def test_stdlib(code: int, args: TestConfig) -> TestResults:
343361
add_files(files, (stdlib / name), args)
344362

345363
if files:
346-
print(f"Testing stdlib ({len(files)} files)...")
347-
print("Running mypy " + " ".join(get_mypy_flags(args, "/tmp/...", testing_stdlib=True)))
348-
this_code = run_mypy(args, [], files, testing_stdlib=True)
364+
print(f"Testing stdlib ({len(files)} files)...", end="", flush=True)
365+
# We don't actually need pip for the stdlib testing
366+
venv_info = VenvInfo(pip_exe="", python_exe=sys.executable)
367+
this_code = run_mypy(args, [], files, venv_info=venv_info, testing_stdlib=True, non_types_dependencies=False)
349368
code = max(code, this_code)
350369

351370
return TestResults(code, len(files))
352371

353372

354-
def test_third_party_stubs(code: int, args: TestConfig) -> TestResults:
373+
_PRINT_LOCK = Lock()
374+
_DISTRIBUTION_TO_VENV_MAPPING: dict[str, VenvInfo] = {}
375+
376+
377+
def setup_venv_for_external_requirements_set(requirements_set: frozenset[str], tempdir: Path) -> tuple[frozenset[str], VenvInfo]:
378+
venv_dir = tempdir / f".venv-{hash(requirements_set)}"
379+
return requirements_set, make_venv(venv_dir)
380+
381+
382+
def install_requirements_for_venv(venv_info: VenvInfo, args: TestConfig, external_requirements: frozenset[str]) -> None:
383+
# Use --no-cache-dir to avoid issues with concurrent read/writes to the cache
384+
pip_command = [venv_info.pip_exe, "install", get_mypy_req(), *sorted(external_requirements), "--no-cache-dir"]
385+
if args.verbose:
386+
with _PRINT_LOCK:
387+
print(colored(f"Running {pip_command}", "blue"))
388+
try:
389+
subprocess.run(pip_command, check=True, capture_output=True, text=True)
390+
except subprocess.CalledProcessError as e:
391+
print(e.stderr)
392+
raise
393+
394+
395+
def setup_virtual_environments(distributions: dict[str, PackageDependencies], args: TestConfig, tempdir: Path) -> None:
396+
# We don't actually need pip if there aren't any external dependencies
397+
no_external_dependencies_venv = VenvInfo(pip_exe="", python_exe=sys.executable)
398+
external_requirements_to_distributions: defaultdict[frozenset[str], list[str]] = defaultdict(list)
399+
num_pkgs_with_external_reqs = 0
400+
401+
for distribution_name, requirements in distributions.items():
402+
if requirements.external_pkgs:
403+
num_pkgs_with_external_reqs += 1
404+
external_requirements = frozenset(requirements.external_pkgs)
405+
external_requirements_to_distributions[external_requirements].append(distribution_name)
406+
else:
407+
_DISTRIBUTION_TO_VENV_MAPPING[distribution_name] = no_external_dependencies_venv
408+
409+
if num_pkgs_with_external_reqs == 0:
410+
if args.verbose:
411+
print(colored("No additional venvs are required to be set up", "blue"))
412+
return
413+
414+
requirements_sets_to_venvs: dict[frozenset[str], VenvInfo] = {}
415+
416+
if args.verbose:
417+
num_venvs = len(external_requirements_to_distributions)
418+
msg = (
419+
f"Setting up {num_venvs} venv{'s' if num_venvs != 1 else ''} "
420+
f"for {num_pkgs_with_external_reqs} "
421+
f"distribution{'s' if num_pkgs_with_external_reqs != 1 else ''}... "
422+
)
423+
print(colored(msg, "blue"), end="", flush=True)
424+
venv_start_time = time.perf_counter()
425+
426+
with concurrent.futures.ThreadPoolExecutor() as executor:
427+
venv_info_futures = [
428+
executor.submit(setup_venv_for_external_requirements_set, requirements_set, tempdir)
429+
for requirements_set in external_requirements_to_distributions
430+
]
431+
for venv_info_future in concurrent.futures.as_completed(venv_info_futures):
432+
requirements_set, venv_info = venv_info_future.result()
433+
requirements_sets_to_venvs[requirements_set] = venv_info
434+
435+
if args.verbose:
436+
venv_elapsed_time = time.perf_counter() - venv_start_time
437+
print(colored(f"took {venv_elapsed_time:.2f} seconds", "blue"))
438+
pip_start_time = time.perf_counter()
439+
440+
# Limit workers to 10 at a time, since this makes network requests
441+
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
442+
pip_install_futures = [
443+
executor.submit(install_requirements_for_venv, venv_info, args, requirements_set)
444+
for requirements_set, venv_info in requirements_sets_to_venvs.items()
445+
]
446+
concurrent.futures.wait(pip_install_futures)
447+
448+
if args.verbose:
449+
pip_elapsed_time = time.perf_counter() - pip_start_time
450+
msg = f"Combined time for installing requirements across all venvs: {pip_elapsed_time:.2f} seconds"
451+
print(colored(msg, "blue"))
452+
453+
for requirements_set, distribution_list in external_requirements_to_distributions.items():
454+
venv_to_use = requirements_sets_to_venvs[requirements_set]
455+
_DISTRIBUTION_TO_VENV_MAPPING.update(dict.fromkeys(distribution_list, venv_to_use))
456+
457+
458+
def test_third_party_stubs(code: int, args: TestConfig, tempdir: Path) -> TestResults:
355459
print("Testing third-party packages...")
356-
print("Running mypy " + " ".join(get_mypy_flags(args, "/tmp/...", testing_stdlib=False)))
357460
files_checked = 0
358461
gitignore_spec = get_gitignore_spec()
462+
distributions_to_check: dict[str, PackageDependencies] = {}
359463

360464
for distribution in sorted(os.listdir("stubs")):
361465
distribution_path = Path("stubs", distribution)
@@ -368,14 +472,25 @@ def test_third_party_stubs(code: int, args: TestConfig) -> TestResults:
368472
or Path("stubs") in args.filter
369473
or any(distribution_path in path.parents for path in args.filter)
370474
):
371-
this_code, checked = test_third_party_distribution(distribution, args)
372-
code = max(code, this_code)
373-
files_checked += checked
475+
distributions_to_check[distribution] = get_recursive_requirements(distribution)
476+
477+
if not _DISTRIBUTION_TO_VENV_MAPPING:
478+
setup_virtual_environments(distributions_to_check, args, tempdir)
479+
480+
assert _DISTRIBUTION_TO_VENV_MAPPING.keys() == distributions_to_check.keys()
481+
482+
for distribution, venv_info in _DISTRIBUTION_TO_VENV_MAPPING.items():
483+
non_types_dependencies = venv_info.python_exe != sys.executable
484+
this_code, checked = test_third_party_distribution(
485+
distribution, args, venv_info=venv_info, non_types_dependencies=non_types_dependencies
486+
)
487+
code = max(code, this_code)
488+
files_checked += checked
374489

375490
return TestResults(code, files_checked)
376491

377492

378-
def test_typeshed(code: int, args: TestConfig) -> TestResults:
493+
def test_typeshed(code: int, args: TestConfig, tempdir: Path) -> TestResults:
379494
print(f"*** Testing Python {args.version} on {args.platform}")
380495
files_checked_this_version = 0
381496
stdlib_dir, stubs_dir = Path("stdlib"), Path("stubs")
@@ -385,7 +500,7 @@ def test_typeshed(code: int, args: TestConfig) -> TestResults:
385500
print()
386501

387502
if stubs_dir in args.filter or any(stubs_dir in path.parents for path in args.filter):
388-
code, third_party_files_checked = test_third_party_stubs(code, args)
503+
code, third_party_files_checked = test_third_party_stubs(code, args, tempdir)
389504
files_checked_this_version += third_party_files_checked
390505
print()
391506

@@ -400,10 +515,12 @@ def main() -> None:
400515
exclude = args.exclude or []
401516
code = 0
402517
total_files_checked = 0
403-
for version, platform in product(versions, platforms):
404-
config = TestConfig(args.verbose, filter, exclude, version, platform)
405-
code, files_checked_this_version = test_typeshed(code, args=config)
406-
total_files_checked += files_checked_this_version
518+
with tempfile.TemporaryDirectory() as td:
519+
td_path = Path(td)
520+
for version, platform in product(versions, platforms):
521+
config = TestConfig(args.verbose, filter, exclude, version, platform)
522+
code, files_checked_this_version = test_typeshed(code, args=config, tempdir=td_path)
523+
total_files_checked += files_checked_this_version
407524
if code:
408525
print_error(f"--- exit status {code}, {total_files_checked} files checked ---")
409526
sys.exit(code)
@@ -417,5 +534,5 @@ def main() -> None:
417534
try:
418535
main()
419536
except KeyboardInterrupt:
420-
print_error("\n\n!!!\nTest aborted due to KeyboardInterrupt\n!!!")
537+
print_error("\n\nTest aborted due to KeyboardInterrupt!")
421538
sys.exit(1)

0 commit comments

Comments
 (0)