Add colour to the output of mypy_test, take 2 (#7879)

Co-authored-by: Sebastian Rittau <srittau@rittau.biz>

.github/workflows/tests.yml

@@ -60,7 +60,7 @@ jobs:
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v2
-      - run: pip install $(grep tomli== requirements-tests.txt) $(grep mypy== requirements-tests.txt)
+      - run: pip install $(grep tomli== requirements-tests.txt) $(grep mypy== requirements-tests.txt) termcolor
       - run: ./tests/mypy_test.py --platform=${{ matrix.platform }} --python-version=${{ matrix.python-version }}
 
   pyright:

tests/colors.py

@@ -0,0 +1,31 @@
+"""
+Helper module so we don't have to install types-termcolor in CI.
+
+This is imported by `mypy_test.py` and `stubtest_third_party.py`.
+"""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+
+    def colored(__str: str, __style: str) -> str:
+        ...
+
+else:
+    try:
+        from termcolor import colored
+    except ImportError:
+
+        def colored(s: str, _: str) -> str:
+            return s
+
+
+def print_error(error: str, end: str = "\n") -> None:
+    error_split = error.split("\n")
+    for line in error_split[:-1]:
+        print(colored(line, "red"))
+    print(colored(error_split[-1], "red"), end=end)
+
+
+def print_success_msg() -> None:
+    print(colored("success", "green"))

tests/mypy_test.py

@@ -19,6 +19,8 @@ import subprocess
 import sys
 import tempfile
 from collections.abc import Iterable
+from contextlib import redirect_stderr, redirect_stdout
+from io import StringIO
 from pathlib import Path
 from typing import TYPE_CHECKING, NamedTuple
 
@@ -28,6 +30,7 @@ if TYPE_CHECKING:
 from typing_extensions import TypeAlias
 
 import tomli
+from colors import colored, print_error, print_success_msg
 
 parser = argparse.ArgumentParser(description="Test runner for typeshed. Patterns are unanchored regexps on the full path.")
 parser.add_argument("-v", "--verbose", action="count", default=0, help="More output")
@@ -165,7 +168,7 @@ def run_mypy(
     try:
         from mypy.api import run as mypy_run
     except ImportError:
-        print("Cannot import mypy. Did you install it?")
+        print_error("Cannot import mypy. Did you install it?")
         sys.exit(1)
 
     with tempfile.NamedTemporaryFile("w+") as temp:
@@ -183,9 +186,24 @@ def run_mypy(
         if args.dry_run:
             exit_code = 0
         else:
-            stdout, stderr, exit_code = mypy_run(mypy_args)
-            print(stdout, end="")
-            print(stderr, end="")
+            stdout_redirect, stderr_redirect = StringIO(), StringIO()
+            with redirect_stdout(stdout_redirect), redirect_stderr(stderr_redirect):
+                returned_stdout, returned_stderr, exit_code = mypy_run(mypy_args)
+
+            if exit_code:
+                print_error("failure\n")
+                captured_stdout = stdout_redirect.getvalue()
+                captured_stderr = stderr_redirect.getvalue()
+                if returned_stderr:
+                    print_error(returned_stderr)
+                if captured_stderr:
+                    print_error(captured_stderr)
+                if returned_stdout:
+                    print_error(returned_stdout)
+                if captured_stdout:
+                    print_error(captured_stdout, end="")
+            else:
+                print_success_msg()
         return exit_code
 
 
@@ -194,8 +212,11 @@ ReturnCode: TypeAlias = int
 
 def run_mypy_as_subprocess(directory: StrPath, flags: Iterable[str]) -> ReturnCode:
     result = subprocess.run([sys.executable, "-m", "mypy", directory, *flags], capture_output=True)
-    print(result.stdout.decode())
-    print(result.stderr.decode())
+    stdout, stderr = result.stdout, result.stderr
+    if stderr:
+        print_error(stderr.decode())
+    if stdout:
+        print_error(stdout.decode())
     return result.returncode
 
 
@@ -234,6 +255,7 @@ def get_mypy_flags(
     if strict:
         flags.append("--strict")
     if test_suite_run:
+        flags.append("--namespace-packages")
         if sys.platform == "win32" or args.platform == "win32":
             flags.extend(["--exclude", "tests/pytype_test.py"])
     else:
@@ -291,10 +313,10 @@ def test_third_party_distribution(distribution: str, major: int, minor: int, arg
     seen_dists: set[str] = set()
     add_third_party_files(distribution, major, files, args, configurations, seen_dists)
 
-    print(f"testing {distribution} ({len(files)} files)...")
+    print(f"testing {distribution} ({len(files)} files)... ", end="")
 
     if not files:
-        print("--- no files found ---")
+        print_error("no files found")
         sys.exit(1)
 
     code = run_mypy(args, configurations, major, minor, files)
@@ -374,6 +396,8 @@ def test_the_test_scripts(code: int, major: int, minor: int, args: argparse.Name
         this_code = 0
     else:
         this_code = run_mypy_as_subprocess("tests", flags)
+    if not this_code:
+        print_success_msg()
     code = max(code, this_code)
     return TestResults(code, num_test_files_to_test)
 
@@ -392,6 +416,8 @@ def test_the_test_cases(code: int, major: int, minor: int, args: argparse.Namesp
         with tempfile.TemporaryDirectory() as td:
             shutil.copytree(Path("test_cases"), Path(td) / "test_cases")
             this_code = run_mypy_as_subprocess(td, flags)
+    if not this_code:
+        print_success_msg()
     code = max(code, this_code)
     return TestResults(code, num_test_case_files)
 
@@ -433,7 +459,7 @@ def main() -> None:
     if args.python_version:
         versions = [v for v in versions if any(("%d.%d" % v).startswith(av) for av in args.python_version)]
         if not versions:
-            print("--- no versions selected ---")
+            print_error("--- no versions selected ---")
             sys.exit(1)
 
     code = 0
@@ -442,13 +468,17 @@ def main() -> None:
         code, files_checked_this_version = test_typeshed(code, major, minor, args)
         total_files_checked += files_checked_this_version
     if code:
-        print(f"--- exit status {code}, {total_files_checked} files checked ---")
+        print_error(f"--- exit status {code}, {total_files_checked} files checked ---")
         sys.exit(code)
     if not total_files_checked:
-        print("--- nothing to do; exit 1 ---")
+        print_error("--- nothing to do; exit 1 ---")
         sys.exit(1)
-    print(f"--- success, {total_files_checked} files checked ---")
+    print(colored(f"--- success, {total_files_checked} files checked ---", "green"))
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except KeyboardInterrupt:
+        print_error("\n\n!!!\nTest aborted due to KeyboardInterrupt\n!!!")
+        sys.exit(1)

tests/stubtest_third_party.py

@@ -10,22 +10,10 @@ import sys
 import tempfile
 import venv
 from pathlib import Path
-from typing import TYPE_CHECKING, NoReturn
+from typing import NoReturn
 
 import tomli
-
-if TYPE_CHECKING:
-
-    def colored(__str: str, __style: str) -> str:
-        ...
-
-else:
-    try:
-        from termcolor import colored
-    except ImportError:
-
-        def colored(s: str, _: str) -> str:
-            return s
+from colors import colored, print_error, print_success_msg
 
 
 @functools.lru_cache()
@@ -38,10 +26,10 @@ def run_stubtest(dist: Path, *, verbose: bool = False) -> bool:
     with open(dist / "METADATA.toml") as f:
         metadata = dict(tomli.loads(f.read()))
 
-    print(f"{dist.name}...", end="")
+    print(f"{dist.name}... ", end="")
 
     if not metadata.get("stubtest", True):
-        print(colored(" skipping", "yellow"))
+        print(colored("skipping", "yellow"))
         return True
 
     with tempfile.TemporaryDirectory() as tmp:
@@ -100,7 +88,7 @@ def run_stubtest(dist: Path, *, verbose: bool = False) -> bool:
         try:
             subprocess.run(stubtest_cmd, env={"MYPYPATH": str(dist), "MYPY_FORCE_COLOR": "1"}, check=True, capture_output=True)
         except subprocess.CalledProcessError as e:
-            print(colored(" fail", "red"))
+            print_error("fail")
             print_commands(dist, pip_cmd, stubtest_cmd)
             print_command_output(e)
 
@@ -120,7 +108,7 @@ def run_stubtest(dist: Path, *, verbose: bool = False) -> bool:
 
             return False
         else:
-            print(colored(" success", "green"))
+            print_success_msg()
 
     if verbose:
         print_commands(dist, pip_cmd, stubtest_cmd)
@@ -136,7 +124,7 @@ def print_commands(dist: Path, pip_cmd: list[str], stubtest_cmd: list[str]) -> N
 
 
 def print_command_failure(message: str, e: subprocess.CalledProcessError) -> None:
-    print(colored(" fail", "red"))
+    print_error("fail")
     print(file=sys.stderr)
     print(message, file=sys.stderr)
     print_command_output(e)