diff --git a/src/west/app/main.py b/src/west/app/main.py index d55fde96..226c5dda 100755 --- a/src/west/app/main.py +++ b/src/west/app/main.py @@ -93,6 +93,7 @@ class EarlyArgs(NamedTuple): version: bool # True if -V was given zephyr_base: str | None # -z argument value verbosity: int # 0 if not given, otherwise counts + color: str | None # --color argument value ('always', 'never', 'auto') command_name: str | None # Other arguments are appended here. @@ -106,16 +107,18 @@ def parse_early_args(argv: list[str]) -> EarlyArgs: version = False zephyr_base = None verbosity = 0 + color = None command_name = None unexpected_arguments = [] expecting_zephyr_base = False + expecting_color = False def consume_more_args(rest): # Handle the 'Vv' portion of 'west -hVv'. - nonlocal help, version, zephyr_base, verbosity - nonlocal expecting_zephyr_base + nonlocal help, version, zephyr_base, verbosity, color + nonlocal expecting_zephyr_base, expecting_color if not rest: return @@ -145,6 +148,13 @@ def consume_more_args(rest): for arg in argv: if expecting_zephyr_base: zephyr_base = arg + expecting_zephyr_base = False + elif expecting_color: + if arg in ('always', 'never', 'auto'): + color = arg + else: + unexpected_arguments.append(f'--color={arg}') + expecting_color = False elif arg.startswith('-h'): help = True consume_more_args(arg[2:]) @@ -170,13 +180,23 @@ def consume_more_args(rest): zephyr_base = arg[3:] else: zephyr_base = arg[2:] + elif arg == '--color': + expecting_color = True + elif arg.startswith('--color='): + color_val = arg[8:] + if color_val in ('always', 'never', 'auto'): + color = color_val + else: + unexpected_arguments.append(arg) elif arg.startswith('-'): unexpected_arguments.append(arg) else: command_name = arg break - return EarlyArgs(help, version, zephyr_base, verbosity, command_name, unexpected_arguments) + return EarlyArgs( + help, version, zephyr_base, verbosity, color, command_name, unexpected_arguments + ) class LogFormatter(logging.Formatter): @@ -222,6 +242,7 @@ def __init__(self): self.subparser_gen = None # an add_subparsers() return value self.cmd = None # west.commands.WestCommand, eventually self.queued_io = [] # I/O hooks we want self.cmd to do + self.color = None # 'always', 'never', 'auto', or None for group, classes in BUILTIN_COMMAND_GROUPS.items(): lst = [cls() for cls in classes] @@ -253,9 +274,8 @@ def run(self, argv): # Use verbosity to determine west API log levels self.setup_west_logging(early_args.verbosity) - # Makes ANSI color escapes work on Windows, and strips them when - # stdout/stderr isn't a terminal - colorama.init() + # Store color mode for later use + self.color = early_args.color # See if we're in a workspace. It's fine if we're not. # Note that this falls back on searching from ZEPHYR_BASE @@ -273,6 +293,13 @@ def run(self, argv): # backwards compatibility. self.config._copy_to_configparser(west.configuration.config) + # Set color preference once. + self.color = self.resolve_color() + self.colorama_init() + + # Propagate app-level color mode to every command instance. + self.queued_io.append(lambda cmd: setattr(cmd, 'color', self.color)) + # Set self.manifest and self.extensions. self.load_manifest() self.load_extension_specs() @@ -576,6 +603,14 @@ def make_parsers(self): help='print the program version and exit', ) + parser.add_argument( + '--color', + choices=['always', 'never', 'auto'], + dest='color', + default=None, + help='when to colorize output (always, never, auto)', + ) + subparser_gen = parser.add_subparsers(metavar='', dest='command') return parser, subparser_gen @@ -739,12 +774,35 @@ def setup_west_logging(self, verbosity): logger.addHandler(LogHandler()) + def resolve_color(self): + if self.color is None: + try: + config_mode = self.config.getboolean('color.ui') + if not config_mode: + return 'never' + else: + return 'auto' + + except ValueError: + return self.config.get('color.ui', 'auto') + else: + return self.color + + def colorama_init(self): + if self.color == 'always': + colorama.init(strip=False) + elif self.color == 'never': + colorama.init(strip=True) + else: + colorama.init() + def run_builtin(self, args, unknown): self.queued_io.append( lambda cmd: cmd.dbg('args namespace:', args, level=Verbosity.DBG_EXTREME) ) self.cmd = self.builtins.get(args.command, self.builtins['help']) adjust_command_verbosity(self.cmd, args) + if self.mle: self.handle_builtin_manifest_load_err(args) for io_hook in self.queued_io: @@ -771,6 +829,7 @@ def run_extension(self, name, argv): args, unknown = west_parser.parse_known_args(argv) adjust_command_verbosity(self.cmd, args) + self.queued_io.append( lambda cmd: cmd.dbg('args namespace:', args, level=Verbosity.DBG_EXTREME) ) diff --git a/src/west/commands.py b/src/west/commands.py index 09f13f8e..dcb1d367 100644 --- a/src/west/commands.py +++ b/src/west/commands.py @@ -167,6 +167,7 @@ def __init__( self.manifest = None self.config = None self._hooks: list[Callable[[WestCommand], None]] = [] + self.color: str | None = None def add_pre_run_hook(self, hook: Callable[['WestCommand'], None]) -> None: '''Add a hook which will be called right before do_run(). @@ -202,6 +203,9 @@ def run( :param config: `west.configuration.Configuration` or ``None``, accessible as ``self.config`` from `WestCommand.do_run` ''' + arg_color = getattr(args, 'color', None) + if arg_color is not None: + self.color = arg_color self.config = config if unknown and not self.accepts_unknown_args: self.parser.error(f'unexpected arguments: {unknown}') @@ -543,8 +547,12 @@ def die(self, *args, exit_code: int = 1) -> NoReturn: @property def color_ui(self) -> bool: - '''Should we colorize output?''' - return self.config.getboolean('color.ui', default=True) if self.has_config else True + if self.color == 'never': + return False + if self.color == 'always': + return True + # Just to represent auto and None return True + return True # # Internal APIs. Not for public consumption. diff --git a/tests/test_commands.py b/tests/test_commands.py index 5f38ca99..da196af3 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -26,6 +26,15 @@ def do_run(self): cmd = WestCommandImpl(name="x", help="y", description="z") + +class DummyConfig: + def __init__(self, color_ui_value): + self.color_ui_value = color_ui_value + + def getboolean(self, _option, default=True): + return self.color_ui_value if self.color_ui_value is not None else default + + TEST_STR = "This is some test string" COL_RED = "\x1b[91m" COL_YELLOW = "\x1b[93m" @@ -120,3 +129,28 @@ def test_die(capsys, test_case): stderr = captured.err assert stderr == exp_err assert stdout == exp_out + + +def test_color_ui_uses_color(): + command = WestCommandImpl(name="x", help="y", description="z") + command.color = 'always' + assert command.color_ui is True + + command.color = 'never' + assert command.color_ui is False + + +def test_color_ui_default_is_enabled_when_mode_is_auto(): + command = WestCommandImpl(name="x", help="y", description="z") + command.color = 'auto' + command.config = DummyConfig(False) + + assert command.color_ui is True + + +def test_color_ui_default_without_config(): + command = WestCommandImpl(name="x", help="y", description="z") + command.color = None + command.config = None + + assert command.color_ui is True diff --git a/tests/test_main.py b/tests/test_main.py index 41fc886c..cad05148 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -5,6 +5,7 @@ import pytest from conftest import cmd, cmd_subprocess +import west.app.main as west_main import west.version @@ -48,3 +49,18 @@ def test_module_run(tmp_path, monkeypatch): # check that that the sys.path was correctly inserted expected_path = Path(__file__).parents[1] / 'src' assert actual_path == [f'{expected_path}', 'initial-path'] + + +@pytest.mark.parametrize( + "argv, expected_color, expected_command, expected_unexpected", + [ + (['--color=always', 'help'], 'always', 'help', []), + (['--color', 'never', 'status'], 'never', 'status', []), + (['--color', 'invalid', 'status'], None, 'status', ['--color=invalid']), + ], +) +def test_parse_early_args_color(argv, expected_color, expected_command, expected_unexpected): + ea = west_main.parse_early_args(argv) + assert ea.color == expected_color + assert ea.command_name == expected_command + assert ea.unexpected_arguments == expected_unexpected