Force colors from mypy always, strip in pants when --no-colors (#16586)

This fixes #16451 by continuing the color half of #16488 (terminal width).

It does this by forcing mypy to always include color escapes, and then stripping them out in pants, based on the [GLOBAL].colors setting.

This seems to only be possible via an undocumented environment variable: `MYPY_FORCE_COLOR`: python/mypy#7771, https://github.com/python/mypy/blob/03638dd670373db0b8f00cc3bcec256d09729d06/mypy/util.py#L543

Doing this also requires setting `TERM=linux` so that curses can find appropriate info in the terminfo database.

[ci skip-rust]
[ci skip-build-wheels]

Author: huonw

src/python/pants/backend/python/typecheck/mypy/rules.py

@@ -46,6 +46,7 @@
 from pants.engine.rules import Get, MultiGet, collect_rules, rule, rule_helper
 from pants.engine.target import CoarsenedTargets, CoarsenedTargetsRequest
 from pants.engine.unions import UnionRule
+from pants.option.global_options import GlobalOptions
 from pants.util.logging import LogLevel
 from pants.util.ordered_set import FrozenOrderedSet, OrderedSet
 from pants.util.strutil import pluralize, shell_quote
@@ -135,6 +136,7 @@ async def mypy_typecheck_partition(
     mkdir: MkdirBinary,
     cp: CpBinary,
     mv: MvBinary,
+    global_options: GlobalOptions,
 ) -> CheckResult:
 
     # MyPy requires 3.5+ to run, but uses the typed-ast library to work with 2.7, 3.4, 3.5, 3.6,
@@ -306,6 +308,13 @@ async def mypy_typecheck_partition(
     env = {
         "PEX_EXTRA_SYS_PATH": ":".join(all_used_source_roots),
         "MYPYPATH": ":".join(all_used_source_roots),
+        # Always emit colors to improve cache hit rates, the results are post-processed to match the
+        # global setting
+        "MYPY_FORCE_COLOR": "1",
+        # Mypy needs to know the terminal so it can use appropriate escape sequences. ansi is a
+        # reasonable lowest common denominator for the sort of escapes mypy uses (NB. TERM=xterm
+        # uses some additional codes that colors.strip_color doesn't remove).
+        "TERM": "ansi",
         # Force a fixed terminal width. This is effectively infinite, disabling mypy's
         # builtin truncation and line wrapping. Terminals do an acceptable job of soft-wrapping
         # diagnostic text and source code is typically already hard-wrapped to a limited width.
@@ -329,7 +338,10 @@ async def mypy_typecheck_partition(
     result = await Get(FallibleProcessResult, Process, process)
     report = await Get(Digest, RemovePrefix(result.output_digest, REPORT_DIR))
     return CheckResult.from_fallible_process_result(
-        result, partition_description=partition.description(), report=report
+        result,
+        partition_description=partition.description(),
+        report=report,
+        strip_formatting=not global_options.colors,
     )
 
 

src/python/pants/backend/python/typecheck/mypy/rules_integration_test.py

@@ -899,4 +899,6 @@ class incredibly_long_type_name_to_force_wrapping_if_mypy_wrapped_error_messages
     assert re.search(
         "error:.*incredibly_long_type_name.*incredibly_long_attribute_name", result[0].stdout
     )
+    # at least one escape sequence that sets text color (red)
+    assert "\033[31m" in result[0].stdout
     assert result[0].report == EMPTY_DIGEST

src/python/pants/core/goals/check.py

@@ -8,6 +8,8 @@
 from dataclasses import dataclass
 from typing import Any, Iterable, cast
 
+import colors
+
 from pants.core.goals.lint import REPORT_DIR as REPORT_DIR  # noqa: F401
 from pants.core.goals.style_request import (
     StyleRequest,
@@ -48,10 +50,13 @@ def from_fallible_process_result(
         *,
         partition_description: str | None = None,
         strip_chroot_path: bool = False,
+        strip_formatting: bool = False,
         report: Digest = EMPTY_DIGEST,
     ) -> CheckResult:
         def prep_output(s: bytes) -> str:
-            return strip_v2_chroot_path(s) if strip_chroot_path else s.decode()
+            chroot = strip_v2_chroot_path(s) if strip_chroot_path else s.decode()
+            formatting = cast(str, colors.strip_color(chroot)) if strip_formatting else chroot
+            return formatting
 
         return CheckResult(
             exit_code=process_result.exit_code,

src/python/pants/core/goals/check_test.py

@@ -8,6 +8,8 @@
 from textwrap import dedent
 from typing import Iterable, Optional, Sequence, Tuple, Type
 
+import pytest
+
 from pants.core.goals.check import (
     Check,
     CheckRequest,
@@ -18,7 +20,9 @@
 )
 from pants.core.util_rules.distdir import DistDir
 from pants.engine.addresses import Address
-from pants.engine.fs import Workspace
+from pants.engine.fs import EMPTY_DIGEST, EMPTY_FILE_DIGEST, Workspace
+from pants.engine.platform import Platform
+from pants.engine.process import FallibleProcessResult, ProcessResultMetadata
 from pants.engine.target import FieldSet, MultipleSourcesField, Target, Targets
 from pants.engine.unions import UnionMembership
 from pants.testutil.option_util import create_options_bootstrapper, create_subsystem
@@ -242,3 +246,34 @@ def test_streaming_output_partitions() -> None:
 
         """
     )
+
+
+@pytest.mark.parametrize(
+    ("strip_chroot_path", "strip_formatting", "expected"),
+    [
+        (False, False, "\033[0;31m/var/pants-sandbox-123/red/path.py\033[0m \033[1mbold\033[0m"),
+        (False, True, "/var/pants-sandbox-123/red/path.py bold"),
+        (True, False, "\033[0;31mred/path.py\033[0m \033[1mbold\033[0m"),
+        (True, True, "red/path.py bold"),
+    ],
+)
+def test_from_fallible_process_result_output_prepping(
+    strip_chroot_path: bool, strip_formatting: bool, expected: str
+) -> None:
+    result = CheckResult.from_fallible_process_result(
+        FallibleProcessResult(
+            exit_code=0,
+            stdout=b"stdout \033[0;31m/var/pants-sandbox-123/red/path.py\033[0m \033[1mbold\033[0m",
+            stdout_digest=EMPTY_FILE_DIGEST,
+            stderr=b"stderr \033[0;31m/var/pants-sandbox-123/red/path.py\033[0m \033[1mbold\033[0m",
+            stderr_digest=EMPTY_FILE_DIGEST,
+            output_digest=EMPTY_DIGEST,
+            platform=Platform.current,
+            metadata=ProcessResultMetadata(0, "ran_locally", 0),
+        ),
+        strip_chroot_path=strip_chroot_path,
+        strip_formatting=strip_formatting,
+    )
+
+    assert result.stdout == "stdout " + expected
+    assert result.stderr == "stderr " + expected