diff --git a/python/scripts/nvbench_compare.py b/python/scripts/nvbench_compare.py index faed762..08f6474 100644 --- a/python/scripts/nvbench_compare.py +++ b/python/scripts/nvbench_compare.py @@ -43,7 +43,7 @@ np: Any = None Fore: Any = None -def load_nvbench_compare_tooling() -> None: +def load_nvbench_compare_tooling(*, load_color: bool = True) -> None: global Fore, np if np is None: @@ -51,7 +51,7 @@ def load_nvbench_compare_tooling() -> None: ToolingDependency("numpy", "numpy", "bulk timing analysis"), tool_name="nvbench-compare", ) - if Fore is None: + if load_color and Fore is None: colorama = require_tooling_dependency( ToolingDependency("colorama", "colorama", "colored status output"), tool_name="nvbench-compare", @@ -3366,7 +3366,7 @@ def main() -> int: return 1 try: - load_nvbench_compare_tooling() + load_nvbench_compare_tooling(load_color=not args.no_color) except MissingToolingDependencyError as exc: print(str(exc), file=sys.stderr) return 1 @@ -3432,7 +3432,10 @@ def main() -> int: print(str(exc), file=sys.stderr) return 1 - warn_fore = Fore.YELLOW if args.ignore_devices else Fore.RED + if args.no_color: + warn_fore = "" + else: + warn_fore = Fore.YELLOW if args.ignore_devices else Fore.RED msg_text = "Device sections do not match" print(colorize(msg_text, warn_fore, Emoji.NONE, args.no_color), end="") print(": ", end="") diff --git a/python/test/test_nvbench_compare.py b/python/test/test_nvbench_compare.py index 747ee50..b03af8d 100644 --- a/python/test/test_nvbench_compare.py +++ b/python/test/test_nvbench_compare.py @@ -197,17 +197,23 @@ def test_compare_tooling_loader_recovers_after_partial_failure( nvbench_compare, "require_tooling_dependency", fake_require_tooling_dependency ) + nvbench_compare.load_nvbench_compare_tooling(load_color=False) + + assert nvbench_compare.np is not None + assert nvbench_compare.Fore is None + assert attempts == [("numpy", "nvbench-compare")] + + attempts.clear() + with pytest.raises(nvbench_compare.MissingToolingDependencyError): nvbench_compare.load_nvbench_compare_tooling() - assert nvbench_compare.np is not None assert nvbench_compare.Fore is None nvbench_compare.load_nvbench_compare_tooling() assert nvbench_compare.Fore is not None assert attempts == [ - ("numpy", "nvbench-compare"), ("colorama", "nvbench-compare"), ("colorama", "nvbench-compare"), ] @@ -2740,6 +2746,39 @@ def test_main_validates_device_filters_before_loading_tooling( assert "--reference-devices" in capsys.readouterr().out +def test_main_no_color_does_not_require_colorama(monkeypatch, nvbench_compare): + calls = [] + + def fake_load_tooling(*, load_color=True): + calls.append(load_color) + nvbench_compare.np = np + if load_color: + raise AssertionError("color tooling should not load") + + monkeypatch.setattr( + nvbench_compare, "load_nvbench_compare_tooling", fake_load_tooling + ) + monkeypatch.setattr( + nvbench_compare.reader, + "read_file", + make_reader_for_roots( + {"devices": [], "benchmarks": []}, + {"devices": [], "benchmarks": []}, + ), + ) + monkeypatch.setattr( + nvbench_compare, "compare_benches", lambda *args, **kwargs: None + ) + monkeypatch.setattr( + sys, + "argv", + ["nvbench_compare", "--no-color", "ref.json", "cmp.json"], + ) + + assert nvbench_compare.main() == 0 + assert calls == [False] + + def test_main_reports_input_read_failures(monkeypatch, capsys, nvbench_compare): def read_file(_): raise OSError("cannot open file")