Properly handle colors in subprocesses

We have to manually specify color flags when calling processes. Added a
few utility functions to specify those.

BUG=b:243079837
TEST=./tools/health-check --color=always -v
./tools/health-check --color=never -v

Change-Id: I9095a888bbff264d15d3dbba41a2800155aa770f
Reviewed-on: https://chromium-review.googlesource.com/c/crosvm/crosvm/+/3880829
Tested-by: Dennis Kempin <denniskempin@google.com>
Commit-Queue: Dennis Kempin <denniskempin@google.com>
Reviewed-by: Daniel Verkamp <dverkamp@chromium.org>

tools/clippy

@@ -14,7 +14,7 @@ from impl.common import CROSVM_ROOT, cwd, run_main, cmd, chdir
 from impl.test_runner import get_workspace_excludes
 from impl.test_target import Triple
 
-clippy = cmd("cargo clippy")
+clippy = cmd("cargo clippy").with_color_flag()
 
 excluded_crates: list[str] = list(get_workspace_excludes(Triple.host_default()))
 features: str = ""

tools/health-check

@@ -11,9 +11,24 @@ from pathlib import Path
 from typing import List
 
 from impl.check_code_hygiene import has_crlf_line_endings
-from impl.common import CROSVM_ROOT, TOOLS_ROOT, argh, chdir, cmd, cwd_context, parallel, run_main
+from impl.common import (
+    CROSVM_ROOT,
+    TOOLS_ROOT,
+    argh,
+    chdir,
+    cmd,
+    cwd_context,
+    parallel,
+    run_main,
+)
 from impl.health_check import Check, CheckContext, run_checks
 
+python = cmd("python3")
+mypy = cmd("mypy").with_color_env("MYPY_FORCE_COLOR")
+black = cmd("black").with_color_arg(always="--color", never="--no-color")
+mdformat = cmd("mdformat")
+lucicfg = cmd("third_party/depot_tools/lucicfg")
+
 
 def check_python_tests(_: CheckContext):
     "No matter which python files have changed, run all available python tests."
@@ -26,13 +41,13 @@ def check_python_tests(_: CheckContext):
 
 def check_python_types(context: CheckContext):
     "Run mypy on all python files to type-check."
-    mypy = cmd("mypy --pretty --color-output").with_env("MYPY_FORCE_COLOR", "1")
-    parallel(*mypy.foreach(context.all_files)).fg(quiet=True)
+    parallel(*mypy("--pretty").foreach(context.all_files)).fg(quiet=True)
 
 
 def check_python_format(context: CheckContext):
-    black = cmd("black", "--check" if not context.fix else "")
-    parallel(*black.foreach(context.modified_files)).fg(quiet=not context.fix)
+    parallel(*black("--check" if not context.fix else None).foreach(context.modified_files)).fg(
+        quiet=not context.fix
+    )
 
 
 def check_crlf_line_endings(_: CheckContext):
@@ -46,13 +61,16 @@ def check_crlf_line_endings(_: CheckContext):
 
 def check_markdown_format(context: CheckContext):
     "Runs mdformat on all markdown files."
-    mdformat = cmd("mdformat --wrap 100", "--check" if not context.fix else "")
-    parallel(*mdformat.foreach(context.modified_files)).fg(quiet=not context.fix)
+    parallel(
+        *mdformat("--wrap 100", "--check" if not context.fix else "").foreach(
+            context.modified_files
+        )
+    ).fg(quiet=not context.fix)
 
 
 def check_rust_clippy(_: CheckContext):
     "Runs clippy on the whole project, no matter which rs files were touched."
-    cmd("./tools/clippy --locked").fg(quiet=True)
+    cmd("./tools/clippy --locked").with_color_flag().fg(quiet=True)
 
 
 def check_rust_format(context: CheckContext):
@@ -61,12 +79,14 @@ def check_rust_format(context: CheckContext):
         rustfmt = cmd(
             cmd("rustup +nightly which rustfmt"),
             "--config imports_granularity=item,group_imports=StdExternalCrate",
-        )
+        ).with_color_flag()
     else:
         rustfmt = cmd(cmd("rustup which rustfmt"))
-    parallel(*rustfmt("--check" if not context.fix else "").foreach(context.modified_files)).fg(
-        quiet=not context.fix
-    )
+    parallel(
+        *rustfmt("--check" if not context.fix else "")
+        .with_color_flag()
+        .foreach(context.modified_files)
+    ).fg(quiet=not context.fix)
 
 
 def check_rust_lockfiles(_: CheckContext):
@@ -126,7 +146,6 @@ def check_copyright_header(context: CheckContext):
 
 def check_infra_configs(context: CheckContext):
     "Validate luci configs by sending them to luci-config."
-    lucicfg = cmd("third_party/depot_tools/lucicfg")
     for file in context.modified_files:
         if context.fix:
             lucicfg("fmt", file).fg()

tools/impl/common.py

@@ -230,6 +230,29 @@ class Command(object):
         cmd.env_vars = {**self.env_vars, "PATH": f"{path_var}:{new_path}"}
         return cmd
 
+    def with_color_arg(
+        self,
+        always: Optional[str] = None,
+        never: Optional[str] = None,
+    ):
+        """Returns a command with an argument added to pass through enabled/disabled colors."""
+        new_cmd = self
+        if color_enabled():
+            if always:
+                new_cmd = new_cmd(always)
+        else:
+            if never:
+                new_cmd = new_cmd(never)
+        return new_cmd
+
+    def with_color_env(self, var_name: str):
+        """Returns a command with an env var added to pass through enabled/disabled colors."""
+        return self.with_env(var_name, "1" if color_enabled() else "0")
+
+    def with_color_flag(self, flag: str = "--color"):
+        """Returns a command with an added --color=always/never/auto flag."""
+        return self.with_color_arg(always=f"{flag}=always", never=f"{flag}=never")
+
     def foreach(self, arguments: Iterable[Any], batch_size: int = 1):
         """
         Yields a new command for each entry in `arguments`.
@@ -615,6 +638,12 @@ def parse_common_args():
 
 def add_common_args(parser: argparse.ArgumentParser):
     "These args are added to all commands."
+    parser.add_argument(
+        "--color",
+        default="auto",
+        choices=("always", "never", "auto"),
+        help="Force enable or disable colors. Defaults to automatic detection.",
+    )
     parser.add_argument(
         "--verbose",
         "-v",
@@ -627,7 +656,7 @@ def add_common_args(parser: argparse.ArgumentParser):
         "-vv",
         action="store_true",
         default=False,
-        help="Print detailed debug information for script developers.",
+        help="Print more debug output",
     )
 
 
@@ -639,6 +668,15 @@ def very_verbose():
     return parse_common_args().very_verbose
 
 
+def color_enabled():
+    color_arg = parse_common_args().color
+    if color_arg == "never":
+        return False
+    if color_arg == "always":
+        return True
+    return sys.stdout.isatty()
+
+
 def all_tracked_files():
     for line in cmd("git ls-files").lines():
         file = Path(line)