Make third party stubtest output more succinct (#7853)

stubtest_third_party.py will now only print a single line when the test
succeeds for a certain distribution, unless the "-v" option is given.
When the test fails or "-v" is given, will still print the full output.

The success status is now colored to make spotting failures easier.

stdout/stderr is now used consistently: The distribution name and
success status is always printed to stdout, all other output goes to
stderr. Running the script with "2>/dev/null" will only show the success
status of the test, one per line.

tests/stubtest_third_party.py

@@ -10,10 +10,23 @@ import sys
 import tempfile
 import venv
 from pathlib import Path
-from typing import NoReturn
+from typing import TYPE_CHECKING, 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
+
 
 @functools.lru_cache()
 def get_mypy_req() -> str:
@@ -21,12 +34,15 @@ def get_mypy_req() -> str:
         return next(line.strip() for line in f if "mypy" in line)
 
 
-def run_stubtest(dist: Path) -> bool:
+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="")
+
     if not metadata.get("stubtest", True):
-        print(f"Skipping stubtest for {dist.name}\n\n")
+        print(colored(" skipping", "yellow"))
+        print()
         return True
 
     with tempfile.TemporaryDirectory() as tmp:
@@ -47,9 +63,7 @@ def run_stubtest(dist: Path) -> bool:
                 pip_cmd = [pip_exe, "install", "-r", str(req_path)]
                 subprocess.run(pip_cmd, check=True, capture_output=True)
             except subprocess.CalledProcessError as e:
-                print(f"Failed to install requirements for {dist.name}", file=sys.stderr)
-                print(e.stdout.decode(), file=sys.stderr)
-                print(e.stderr.decode(), file=sys.stderr)
+                print_command_failure("Failed to install requirements", e)
                 return False
 
         # We need stubtest to be able to import the package, so install mypy into the venv
@@ -58,18 +72,15 @@ def run_stubtest(dist: Path) -> bool:
         dists_to_install = [dist_req, get_mypy_req()]
         dists_to_install.extend(metadata.get("requires", []))
         pip_cmd = [pip_exe, "install"] + dists_to_install
-        print(" ".join(pip_cmd), file=sys.stderr)
         try:
             subprocess.run(pip_cmd, check=True, capture_output=True)
         except subprocess.CalledProcessError as e:
-            print(f"Failed to install {dist.name}", file=sys.stderr)
-            print(e.stdout.decode(), file=sys.stderr)
-            print(e.stderr.decode(), file=sys.stderr)
+            print_command_failure("Failed to install", e)
             return False
 
         packages_to_check = [d.name for d in dist.iterdir() if d.is_dir() and d.name.isidentifier()]
         modules_to_check = [d.stem for d in dist.iterdir() if d.is_file() and d.suffix == ".pyi"]
-        cmd = [
+        stubtest_cmd = [
             python_exe,
             "-m",
             "mypy.stubtest",
@@ -85,35 +96,62 @@ def run_stubtest(dist: Path) -> bool:
         ]
         allowlist_path = dist / "@tests/stubtest_allowlist.txt"
         if allowlist_path.exists():
-            cmd.extend(["--allowlist", str(allowlist_path)])
+            stubtest_cmd.extend(["--allowlist", str(allowlist_path)])
 
         try:
-            print(f"MYPYPATH={dist}", " ".join(cmd), file=sys.stderr)
-            subprocess.run(cmd, env={"MYPYPATH": str(dist), "MYPY_FORCE_COLOR": "1"}, check=True)
-        except subprocess.CalledProcessError:
-            print(f"stubtest failed for {dist.name}", file=sys.stderr)
-            print("\n\n", file=sys.stderr)
+            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_commands(dist, pip_cmd, stubtest_cmd)
+            print_command_output(e)
 
             print("Ran with the following environment:", file=sys.stderr)
-            subprocess.run([pip_exe, "freeze"])
+            ret = subprocess.run([pip_exe, "freeze"], capture_output=True)
+            print_command_output(ret)
 
             if allowlist_path.exists():
                 print(
                     f'To fix "unused allowlist" errors, remove the corresponding entries from {allowlist_path}', file=sys.stderr
                 )
+                print(file=sys.stderr)
             else:
                 print(f"Re-running stubtest with --generate-allowlist.\nAdd the following to {allowlist_path}:", file=sys.stderr)
-                subprocess.run(cmd + ["--generate-allowlist"], env={"MYPYPATH": str(dist)})
-                print("\n\n", file=sys.stderr)
+                ret = subprocess.run(stubtest_cmd + ["--generate-allowlist"], env={"MYPYPATH": str(dist)}, capture_output=True)
+                print_command_output(ret)
+
             return False
         else:
-            print(f"stubtest succeeded for {dist.name}", file=sys.stderr)
-        print("\n\n", file=sys.stderr)
+            print(colored(" success", "green"))
+
+    if verbose:
+        print_commands(dist, pip_cmd, stubtest_cmd)
+
     return True
 
 
+def print_commands(dist: Path, pip_cmd: list[str], stubtest_cmd: list[str]) -> None:
+    print(file=sys.stderr)
+    print(" ".join(pip_cmd), file=sys.stderr)
+    print(f"MYPYPATH={dist}", " ".join(stubtest_cmd), file=sys.stderr)
+    print(file=sys.stderr)
+
+
+def print_command_failure(message: str, e: subprocess.CalledProcessError) -> None:
+    print(colored(" fail", "red"))
+    print(file=sys.stderr)
+    print(message, file=sys.stderr)
+    print_command_output(e)
+
+
+def print_command_output(e: subprocess.CalledProcessError | subprocess.CompletedProcess[bytes]) -> None:
+    print(e.stdout.decode(), end="", file=sys.stderr)
+    print(e.stderr.decode(), end="", file=sys.stderr)
+    print(file=sys.stderr)
+
+
 def main() -> NoReturn:
     parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--verbose", action="store_true", help="verbose output")
     parser.add_argument("--num-shards", type=int, default=1)
     parser.add_argument("--shard-index", type=int, default=0)
     parser.add_argument("dists", metavar="DISTRIBUTION", type=str, nargs=argparse.ZERO_OR_MORE)
@@ -129,10 +167,13 @@ def main() -> NoReturn:
     for i, dist in enumerate(dists):
         if i % args.num_shards != args.shard_index:
             continue
-        if not run_stubtest(dist):
+        if not run_stubtest(dist, verbose=args.verbose):
             result = 1
     sys.exit(result)
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except KeyboardInterrupt:
+        pass