From 2a1a3703c9194b55583bdb6b4318c7c2b625b100 Mon Sep 17 00:00:00 2001 From: zhenchaoni Date: Mon, 29 Jun 2026 13:46:58 +0800 Subject: [PATCH] Implement no-color option --- docs/commands/overview.md | 4 +++ src/winml/modelkit/cli.py | 3 +- src/winml/modelkit/commands/analyze.py | 1 + src/winml/modelkit/commands/build.py | 1 + src/winml/modelkit/commands/compile.py | 1 + src/winml/modelkit/commands/config.py | 1 + src/winml/modelkit/commands/eval.py | 1 + src/winml/modelkit/commands/export.py | 1 + src/winml/modelkit/commands/inspect.py | 1 + src/winml/modelkit/commands/optimize.py | 1 + src/winml/modelkit/commands/perf.py | 1 + src/winml/modelkit/commands/quantize.py | 1 + src/winml/modelkit/commands/sys.py | 1 + src/winml/modelkit/utils/cli.py | 29 ++++++++++++++++++ tests/unit/utils/test_cli.py | 40 +++++++++++++++++++++++++ 15 files changed, 86 insertions(+), 1 deletion(-) diff --git a/docs/commands/overview.md b/docs/commands/overview.md index 473a5bb73..0e9374e6e 100644 --- a/docs/commands/overview.md +++ b/docs/commands/overview.md @@ -63,6 +63,10 @@ differ per command (e.g., `-p` is a short form for `--precision` only on `config` and `quantize`); check the **Flags** section of each command page rather than assuming they transfer. +`--no-color` is accepted by every command and disables colored output for +that invocation. Color is also auto-disabled when `NO_COLOR=1` or `CI=true` +is set in the environment. + ## See also - [How winml-cli Works](../concepts/how-it-works.md) — end-to-end pipeline overview diff --git a/src/winml/modelkit/cli.py b/src/winml/modelkit/cli.py index d910242c7..fe36e905c 100644 --- a/src/winml/modelkit/cli.py +++ b/src/winml/modelkit/cli.py @@ -34,7 +34,7 @@ from . import __version__ from .telemetry import ActionGroup from .telemetry import telemetry as _telemetry_mod -from .utils.cli import verbosity_options +from .utils.cli import no_color_option, verbosity_options from .utils.logging import configure_logging, flush_ort_startup_logs @@ -249,6 +249,7 @@ def format_commands(self, ctx: click.Context, formatter: click.HelpFormatter) -> ) @click.version_option(version=__version__, prog_name="winml") @verbosity_options() +@no_color_option() @click.option( "--debug", is_flag=True, diff --git a/src/winml/modelkit/commands/analyze.py b/src/winml/modelkit/commands/analyze.py index 741ade720..28df766da 100644 --- a/src/winml/modelkit/commands/analyze.py +++ b/src/winml/modelkit/commands/analyze.py @@ -736,6 +736,7 @@ def _build_runtime_debug_output_path(model_path: Path, ep_name: str, device_name ), ) @cli_utils.verbosity_options() +@cli_utils.no_color_option() @cli_utils.build_config_option() @cli_utils.output_option("Save JSON output to file") @cli_utils.overwrite_option(optional_message="Applies to both --output and --optim-config.") diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index 4cdfd1f04..2c748d527 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -475,6 +475,7 @@ def _validate_loader_tasks_for_model( optional_message="Trust remote code for custom model architectures (e.g., Mu2)." ) @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def build( ctx: click.Context, diff --git a/src/winml/modelkit/commands/compile.py b/src/winml/modelkit/commands/compile.py index 573093876..1a4070032 100644 --- a/src/winml/modelkit/commands/compile.py +++ b/src/winml/modelkit/commands/compile.py @@ -100,6 +100,7 @@ ) @cli_utils.build_config_option() @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def compile( ctx: click.Context, diff --git a/src/winml/modelkit/commands/config.py b/src/winml/modelkit/commands/config.py index 9414e6e8b..c3e480312 100644 --- a/src/winml/modelkit/commands/config.py +++ b/src/winml/modelkit/commands/config.py @@ -128,6 +128,7 @@ def _apply_stage_overrides(cfg: Any, *, no_quant: bool, no_compile: bool) -> Non ) @cli_utils.trust_remote_code_option() @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def config( ctx: click.Context, diff --git a/src/winml/modelkit/commands/eval.py b/src/winml/modelkit/commands/eval.py index e43e31e6b..e1295784c 100644 --- a/src/winml/modelkit/commands/eval.py +++ b/src/winml/modelkit/commands/eval.py @@ -167,6 +167,7 @@ @cli_utils.format_option() @cli_utils.build_config_option() @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def eval( ctx: click.Context, diff --git a/src/winml/modelkit/commands/export.py b/src/winml/modelkit/commands/export.py index ee2b93bfa..b7eb7cf67 100644 --- a/src/winml/modelkit/commands/export.py +++ b/src/winml/modelkit/commands/export.py @@ -132,6 +132,7 @@ def _delete_onnx_with_external_data(onnx_path: Path) -> None: ) @cli_utils.build_config_option() @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def export( ctx: click.Context, diff --git a/src/winml/modelkit/commands/inspect.py b/src/winml/modelkit/commands/inspect.py index 14d7cf773..22e664fb4 100644 --- a/src/winml/modelkit/commands/inspect.py +++ b/src/winml/modelkit/commands/inspect.py @@ -124,6 +124,7 @@ def _looks_like_local_path(model_id: str) -> bool: help="Override model class (e.g., BertForMaskedLM) — can be used without --model", ) @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def inspect( ctx: click.Context, diff --git a/src/winml/modelkit/commands/optimize.py b/src/winml/modelkit/commands/optimize.py index 02010bb35..e89dc5366 100644 --- a/src/winml/modelkit/commands/optimize.py +++ b/src/winml/modelkit/commands/optimize.py @@ -182,6 +182,7 @@ def capability_options(func: F) -> F: help="Configuration file (YAML/JSON)", ) @cli_utils.verbosity_options() +@cli_utils.no_color_option() @capability_options @click.pass_context # type: ignore[arg-type] # capability_options widens the signature; click stubs want positional-only ctx but we keep it keyword-callable for back-compat def optimize( diff --git a/src/winml/modelkit/commands/perf.py b/src/winml/modelkit/commands/perf.py index 6c92461d5..0f83c871c 100644 --- a/src/winml/modelkit/commands/perf.py +++ b/src/winml/modelkit/commands/perf.py @@ -1544,6 +1544,7 @@ def _run_simple_loop( @cli_utils.format_option() @cli_utils.build_config_option() @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def perf( ctx: click.Context, diff --git a/src/winml/modelkit/commands/quantize.py b/src/winml/modelkit/commands/quantize.py index 079b37ea8..29f528205 100644 --- a/src/winml/modelkit/commands/quantize.py +++ b/src/winml/modelkit/commands/quantize.py @@ -105,6 +105,7 @@ ) @cli_utils.build_config_option() @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def quantize( ctx: click.Context, diff --git a/src/winml/modelkit/commands/sys.py b/src/winml/modelkit/commands/sys.py index 96023ec7a..ed107023f 100644 --- a/src/winml/modelkit/commands/sys.py +++ b/src/winml/modelkit/commands/sys.py @@ -680,6 +680,7 @@ def _output_ep_text(eps: list[dict[str, Any]]) -> None: help="List available execution providers", ) @cli_utils.verbosity_options() +@cli_utils.no_color_option() @click.pass_context def sysinfo( ctx: click.Context, diff --git a/src/winml/modelkit/utils/cli.py b/src/winml/modelkit/utils/cli.py index c62ede36b..6f6648eb3 100644 --- a/src/winml/modelkit/utils/cli.py +++ b/src/winml/modelkit/utils/cli.py @@ -7,6 +7,7 @@ from __future__ import annotations import json +import os from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypeVar @@ -514,6 +515,34 @@ def decorator(f: F) -> F: return decorator +def no_color_option() -> Callable[[F], F]: + """Add a ``--no-color`` flag that disables colored output. + + Rich honors the ``NO_COLOR`` environment variable for every Console, so the + flag's callback just sets ``NO_COLOR=1`` for the remainder of the run — this + covers all consoles regardless of how they are constructed and matches the + existing ``NO_COLOR=1`` / ``CI=true`` environment behavior. The change lives + only in the current process, so the next invocation is colored again. + + Returns: + Decorator function adding the ``--no-color`` flag (no exposed param). + """ + + def _disable_color(ctx: click.Context, param: click.Parameter, value: bool) -> bool: + if value: + os.environ["NO_COLOR"] = "1" + return value + + return click.option( + "--no-color", + is_flag=True, + default=False, + expose_value=False, + callback=_disable_color, + help="Disable colored output (also via NO_COLOR=1 or CI=true).", + ) + + def resolve_verbosity(ctx: click.Context, verbose: int, quiet: bool) -> tuple[int, bool]: """Merge subcommand ``--verbose``/``--quiet`` with the parent group's values. diff --git a/tests/unit/utils/test_cli.py b/tests/unit/utils/test_cli.py index 5eb2d08a1..6955c7c41 100644 --- a/tests/unit/utils/test_cli.py +++ b/tests/unit/utils/test_cli.py @@ -6,6 +6,7 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING import click @@ -18,6 +19,7 @@ guard_output, ignored_build_flags_warning, max_optim_iterations_option, + no_color_option, optimize_option, overwrite_option, parse_ep_options, @@ -74,6 +76,44 @@ def test_empty_key_raises(self) -> None: parse_ep_options(("=value",)) +class TestNoColorOption: + """Tests for the shared no_color_option() decorator.""" + + @staticmethod + def _make_cmd() -> click.Command: + @click.command() + @no_color_option() + def cmd() -> None: + click.echo("NO_COLOR" in os.environ) + + return cmd + + def test_no_flag_leaves_env_unset(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Without --no-color, NO_COLOR is not set by the callback.""" + monkeypatch.delenv("NO_COLOR", raising=False) + result = CliRunner().invoke(self._make_cmd(), []) + assert result.exit_code == 0 + assert result.output.strip() == "False" + + def test_flag_sets_no_color_env(self, monkeypatch: pytest.MonkeyPatch) -> None: + """--no-color sets NO_COLOR=1 so Rich disables color for the run.""" + monkeypatch.delenv("NO_COLOR", raising=False) + result = CliRunner().invoke(self._make_cmd(), ["--no-color"]) + assert result.exit_code == 0 + assert result.output.strip() == "True" + + def test_flag_not_exposed_as_param(self) -> None: + """expose_value=False: the command takes no extra parameter.""" + result = CliRunner().invoke(self._make_cmd(), ["--no-color"]) + assert result.exit_code == 0 + + def test_help_documents_env_vars(self) -> None: + """Help mentions the NO_COLOR / CI environment fallbacks.""" + result = CliRunner().invoke(self._make_cmd(), ["--help"]) + assert "--no-color" in result.output + assert "NO_COLOR" in result.output + + class TestPrecisionOption: """Tests for the shared precision_option() decorator."""