Py tooling: fix Windows terminal color output and handle color flags in compare.py (fixes #642)

AUTHORS

@@ -41,6 +41,7 @@ Haihan Jiang <haihanj99@gmail.com>
 Henrique Bucher <hbucher@gmail.com>
 International Business Machines Corporation
 Ismael Jimenez Martinez <ismael.jimenez.martinez@gmail.com>
+jdymitarai <o10040115@gmail.com>
 Jern-Kuan Leong <jernkuan@gmail.com>
 JianXiong Zhou <zhoujianxiong2@gmail.com>
 Joao Paulo Magalhaes <joaoppmagalhaes@gmail.com>

CONTRIBUTORS

@@ -62,6 +62,7 @@ Hannes Hauswedell <h2@fsfe.org>
 Henrique Bucher <hbucher@gmail.com>
 Ismael Jimenez Martinez <ismael.jimenez.martinez@gmail.com>
 Iakov Sergeev <yahontu@gmail.com>
+jdymitarai <o10040115@gmail.com>
 Jern-Kuan Leong <jernkuan@gmail.com>
 JianXiong Zhou <zhoujianxiong2@gmail.com>
 Joao Paulo Magalhaes <joaoppmagalhaes@gmail.com>

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()

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",