python · hamdanal · Sep 30, 2023 · Sep 30, 2023 · Oct 1, 2023 · Oct 1, 2023
diff --git a/tests/pytype_test.py b/tests/pytype_test.py
@@ -19,12 +19,16 @@
 import inspect
 import os
 import sys
+import textwrap
 import traceback
+from collections import defaultdict
 from collections.abc import Iterable, Sequence
+from pathlib import Path
 
 from packaging.requirements import Requirement
 
-from parse_metadata import read_dependencies
+from parse_metadata import read_dependencies, read_metadata
+from utils import PYTHON_VERSION, chdir, colored
 
 if sys.platform == "win32":
     print("pytype does not support Windows.", file=sys.stderr)
@@ -45,12 +49,16 @@
 def main() -> None:
     args = create_parser().parse_args()
     typeshed_location = args.typeshed_location or os.getcwd()
-    subdir_paths = [os.path.join(typeshed_location, d) for d in TYPESHED_SUBDIRS]
-    check_subdirs_discoverable(subdir_paths)
     old_typeshed_home = os.environ.get(TYPESHED_HOME)
     os.environ[TYPESHED_HOME] = typeshed_location
-    files_to_test = determine_files_to_test(paths=args.files or subdir_paths)
-    run_all_tests(files_to_test=files_to_test, print_stderr=args.print_stderr, dry_run=args.dry_run)
+    print(f"Testing files with pytype on Python {PYTHON_VERSION}...")
+    with chdir(typeshed_location):
+        check_subdirs_discoverable(TYPESHED_SUBDIRS)
+        files_to_test = determine_files_to_test(paths=args.files or TYPESHED_SUBDIRS)
+        if not files_to_test:
+            print(colored("Nothing to do; exit 1.", "red"))
+            sys.exit(1)
+        run_all_tests(files_to_test=files_to_test, print_stderr=args.print_stderr, dry_run=args.dry_run)
     if old_typeshed_home is None:
         del os.environ[TYPESHED_HOME]
     else:
@@ -122,11 +130,41 @@ def check_subdirs_discoverable(subdir_paths: list[str]) -> None:
             raise SystemExit(f"Cannot find typeshed subdir at {p} (specify parent dir via --typeshed-location)")
 
 
+def classify_files(paths: Sequence[str]) -> tuple[list[str], defaultdict[str, list[str]]]:
+    """Classify files into stdlib and stubs by distribution."""
+    stdlib: list[str] = []
+    stubs: defaultdict[str, list[str]] = defaultdict(list)
+    stubs_path = Path("stubs")
+    stubs_absolute_path = stubs_path.resolve()
+    for path_s in paths:
+        path = Path(path_s).resolve()
+        if path.samefile(stubs_absolute_path):
+            # All stubs, classify by distribution for version checking later.
+            for subdir in stubs_path.iterdir():
+                stubs[subdir.name].append(str(subdir))
+        elif path.is_relative_to(stubs_absolute_path):
+            # A single stub directory or file.
+            distribution = path.relative_to(stubs_absolute_path).parts[0]
+            stubs[distribution].append(path_s)
+        else:
+            stdlib.append(path_s)
+    return stdlib, stubs
+
+
 def determine_files_to_test(*, paths: Sequence[str]) -> list[str]:
     """Determine all files to test, checking if it's in the exclude list and which Python versions to use.
 
     Returns a list of pairs of the file path and Python version as an int."""
-    filenames = find_stubs_in_paths(paths)
+    stdlib, stubs = classify_files(paths)
+    paths_to_test = list(stdlib)
+    for pkg, pkg_paths in stubs.items():
+        requires_python = read_metadata(pkg).requires_python
+        if not requires_python.contains(PYTHON_VERSION):
+            msg = f"skipping {pkg!r} (requires Python {requires_python}; test is being run using Python {PYTHON_VERSION})"
+            print(colored(msg, "yellow"))
+            continue
+        paths_to_test.extend(pkg_paths)
+    filenames = find_stubs_in_paths(paths_to_test)
     ts = typeshed.Typeshed()
     skipped = set(ts.read_blacklist())
     files = []
@@ -208,29 +246,40 @@ def run_all_tests(*, files_to_test: Sequence[str], print_stderr: bool, dry_run:
     errors = 0
     total_tests = len(files_to_test)
     missing_modules = get_missing_modules(files_to_test)
-    print("Testing files with pytype...")
-    for i, f in enumerate(files_to_test):
-        python_version = "{0.major}.{0.minor}".format(sys.version_info)
+    for runs, f in enumerate(files_to_test, start=1):
         if dry_run:
             stderr = None
         else:
-            stderr = run_pytype(filename=f, python_version=python_version, missing_modules=missing_modules)
+            stderr = run_pytype(filename=f, python_version=PYTHON_VERSION, missing_modules=missing_modules)
         if stderr:
-            if print_stderr:
-                print(f"\n{stderr}")
             errors += 1
+            test_file = f"{_get_relative(f)}:"
+            if print_stderr:
+                print(colored(test_file, "red"))
+                print(f"{textwrap.indent(stderr, '  ')}")
             stacktrace_final_line = stderr.rstrip().rsplit("\n", 1)[-1]
-            bad.append((_get_relative(f), python_version, stacktrace_final_line))
+            bad.append((test_file, stacktrace_final_line))
 
-        runs = i + 1
         if runs % 25 == 0:
-            print(f"  {runs:3d}/{total_tests:d} with {errors:3d} errors")
+            color = "red" if errors else "green"
+            print(colored(f"  {runs:4d}/{total_tests:d} with {errors:4d} errors", color))
 
-    print(f"Ran pytype with {total_tests:d} pyis, got {errors:d} errors.")
-    for f, v, err in bad:
-        print(f"\n{f} ({v}): {err}")
+    for test, err in bad:
+        print(colored(test, "red"), err)
+
+    file_plural = "file" if total_tests == 1 else "files"
+    error_plural = "error" if errors == 1 else "errors"
+    msg = f"\nRan pytype with {total_tests:d} pyi {file_plural}, got {errors:d} {error_plural}."
     if errors:
-        raise SystemExit("\nRun again with --print-stderr to get the full stacktrace.")
+        color = "red"
+        code = 1
+        if not print_stderr:
+            msg += "\nRun again with --print-stderr to get the full stacktrace."
+    else:
+        color = "green"
+        code = 0
+    print(colored(msg, color))
+    sys.exit(code)
 
 
 if __name__ == "__main__":

diff --git a/tests/utils.py b/tests/utils.py
@@ -7,6 +7,8 @@
 import subprocess
 import sys
 import venv
+from collections.abc import Iterator
+from contextlib import contextmanager
 from functools import lru_cache
 from pathlib import Path
 from typing import Any, Final, NamedTuple
@@ -138,3 +140,16 @@ def spec_matches_path(spec: pathspec.PathSpec, path: Path) -> bool:
     if path.is_dir():
         normalized_path += "/"
     return spec.match_file(normalized_path)
+
+
+# ====================================================================
+# Similar to `contextlib.chdir` on Python 3.11+
+# ====================================================================
+@contextmanager
+def chdir(path: str | os.PathLike[str]) -> Iterator[None]:
+    old_cwd = os.getcwd()
+    try:
+        os.chdir(path)
+        yield
+    finally:
+        os.chdir(old_cwd)