From 80784e65df009c1c4ba8f9f92f0bc6970c362e8a Mon Sep 17 00:00:00 2001 From: jdymitarai Date: Thu, 10 Sep 2026 09:48:41 +0800 Subject: [PATCH] Py tooling: fix Windows terminal color output and handle color flags in compare.py (fixes #642) * Enable virtual terminal processing on Windows console handles when supported so ANSI escape codes are rendered natively instead of printing literal '^[[92m' sequences. * Fall back to disabled colors when stdout is not a TTY or when the standard NO_COLOR environment variable is present. * Support --color and --no-color on the root parser and correctly extract --color, --no-color, and --benchmark_color=false|true|auto if provided after subcommands. * Prevent false warnings in check_inputs when --benchmark_color is provided with JSON inputs. * Fix 'import util' in tools/gbench/report.py test to 'from gbench import util'. * Add unit tests covering color flag parsing, option extraction, and NO_COLOR detection. --- AUTHORS | 1 + CONTRIBUTORS | 1 + tools/compare.py | 151 ++++++++++++++++++++++++++++++++++++++++- tools/gbench/report.py | 2 +- 4 files changed, 152 insertions(+), 3 deletions(-) diff --git a/AUTHORS b/AUTHORS index 64adf363ab..65cad404a0 100644 --- a/AUTHORS +++ b/AUTHORS @@ -41,6 +41,7 @@ Haihan Jiang Henrique Bucher International Business Machines Corporation Ismael Jimenez Martinez +jdymitarai Jern-Kuan Leong JianXiong Zhou Joao Paulo Magalhaes diff --git a/CONTRIBUTORS b/CONTRIBUTORS index 91f52474f2..a77788f842 100644 --- a/CONTRIBUTORS +++ b/CONTRIBUTORS @@ -62,6 +62,7 @@ Hannes Hauswedell Henrique Bucher Ismael Jimenez Martinez Iakov Sergeev +jdymitarai Jern-Kuan Leong JianXiong Zhou Joao Paulo Magalhaes diff --git a/tools/compare.py b/tools/compare.py index 1a656345c2..5daff15dea 100755 --- a/tools/compare.py +++ b/tools/compare.py @@ -41,6 +41,8 @@ def check_inputs(in1, in2, flags): # When both sides are JSON the only supported flag is # --benchmark_filter= for flag in util.remove_benchmark_flags("--benchmark_filter=", flags): + if flag.startswith("--benchmark_color="): + continue print( "WARNING: passing %s has no effect since both " "inputs are JSON" % flag @@ -56,6 +58,77 @@ def check_inputs(in1, in2, flags): sys.exit(1) +def enable_virtual_terminal_processing(): + """ + On Windows, enable ENABLE_VIRTUAL_TERMINAL_PROCESSING on the console + output handle to allow ANSI escape sequences to be rendered natively. + Returns True if VT processing is enabled or on non-Windows; False otherwise. + """ + if sys.platform == "win32": + try: + import ctypes + + kernel32 = ctypes.windll.kernel32 + handle = kernel32.GetStdHandle(-11) # STD_OUTPUT_HANDLE + mode = ctypes.c_ulong() + if kernel32.GetConsoleMode(handle, ctypes.byref(mode)): + # ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004 + return bool( + kernel32.SetConsoleMode(handle, mode.value | 0x0004) + ) + return False + except Exception: + return False + return True + + +def should_use_color(): + """ + Determine whether terminal color output should be enabled by default. + Follows the NO_COLOR specification (https://no-color.org/), verifies that + stdout is a TTY, and checks that the Windows console supports VT processing. + """ + if "NO_COLOR" in os.environ: + return False + if not hasattr(sys.stdout, "isatty") or not sys.stdout.isatty(): + return False + if sys.platform == "win32": + return enable_virtual_terminal_processing() + return os.environ.get("TERM", "") != "dumb" + + +def resolve_color(color_flag, benchmark_options): + """Determine whether to use color and strip internal color options. + + Returns (use_color, remaining_benchmark_options). + """ + remaining = [] + color = color_flag + for opt in benchmark_options: + if opt == "--no-color": + color = False + elif opt == "--color": + color = True + elif opt.startswith("--benchmark_color="): + val = opt.split("=", 1)[1].lower() + if val in ("false", "0", "no"): + color = False + elif val in ("true", "1", "yes"): + color = True + elif val == "auto": + color = None + remaining.append(opt) + else: + remaining.append(opt) + + if color is None: + color = should_use_color() + elif color: + enable_virtual_terminal_processing() + + return color, remaining + + def create_parser(): parser = ArgumentParser( description="versatile benchmark output compare tool" @@ -73,10 +146,17 @@ def create_parser(): "Internally, all the actual runs are still used, e.g. for U test.", ) - parser.add_argument( + color_group = parser.add_mutually_exclusive_group() + color_group.add_argument( + "--color", + dest="color", + default=None, + action="store_true", + help="Always use colors in the terminal output", + ) + color_group.add_argument( "--no-color", dest="color", - default=True, action="store_false", help="Do not use colors in the terminal output", ) @@ -254,6 +334,7 @@ def main(): exit(1) assert not unknown_args benchmark_options = args.benchmark_options + args.color, benchmark_options = resolve_color(args.color, benchmark_options) if args.mode == "benchmarks": test_baseline = args.test_baseline[0].name @@ -534,6 +615,72 @@ def test_benchmarksfiltered_with_remainder_after_doubleminus(self): self.assertEqual(parsed.filter_contender[0], "e") self.assertEqual(parsed.benchmark_options[0], "g") + def test_benchmarks_color_flag(self): + parsed = self.parser.parse_args( + ["--color", "benchmarks", self.testInput0, self.testInput1] + ) + self.assertTrue(parsed.color) + + def test_benchmarks_no_color_flag(self): + parsed = self.parser.parse_args( + ["--no-color", "benchmarks", self.testInput0, self.testInput1] + ) + self.assertFalse(parsed.color) + + def test_benchmarks_color_default_none(self): + parsed = self.parser.parse_args( + ["benchmarks", self.testInput0, self.testInput1] + ) + self.assertIsNone(parsed.color) + + +class TestColorResolution(unittest.TestCase): + def test_resolve_color_root_flags(self): + color, remaining = resolve_color(True, ["opt1"]) + self.assertTrue(color) + self.assertEqual(remaining, ["opt1"]) + + color, remaining = resolve_color(False, ["opt1"]) + self.assertFalse(color) + self.assertEqual(remaining, ["opt1"]) + + def test_resolve_color_remainder_flags(self): + color, remaining = resolve_color(None, ["--no-color", "opt1"]) + self.assertFalse(color) + self.assertEqual(remaining, ["opt1"]) + + color, remaining = resolve_color(None, ["--color", "opt1"]) + self.assertTrue(color) + self.assertEqual(remaining, ["opt1"]) + + def test_resolve_color_benchmark_color_flag(self): + color, remaining = resolve_color(None, ["--benchmark_color=false"]) + self.assertFalse(color) + self.assertEqual(remaining, ["--benchmark_color=false"]) + + color, remaining = resolve_color(None, ["--benchmark_color=0"]) + self.assertFalse(color) + self.assertEqual(remaining, ["--benchmark_color=0"]) + + color, remaining = resolve_color(None, ["--benchmark_color=true"]) + self.assertTrue(color) + self.assertEqual(remaining, ["--benchmark_color=true"]) + + color, remaining = resolve_color(None, ["--benchmark_color=1"]) + self.assertTrue(color) + self.assertEqual(remaining, ["--benchmark_color=1"]) + + def test_should_use_color_no_color_env(self): + orig = os.environ.get("NO_COLOR") + try: + os.environ["NO_COLOR"] = "1" + self.assertFalse(should_use_color()) + finally: + if orig is None: + os.environ.pop("NO_COLOR", None) + else: + os.environ["NO_COLOR"] = orig + if __name__ == "__main__": # unittest.main() diff --git a/tools/gbench/report.py b/tools/gbench/report.py index 16b674ab09..51a61dbedf 100644 --- a/tools/gbench/report.py +++ b/tools/gbench/report.py @@ -1489,7 +1489,7 @@ def load_result(): cls.json = load_result() def test_json_diff_report_pretty_printing(self): - import util + from gbench import util expected_names = [ "99 family 0 instance 0 repetition 0",