From 9716d3386f030f89c525f18b84d6d022d3da804f Mon Sep 17 00:00:00 2001 From: Alexander Goscinski Date: Tue, 18 Jun 2024 15:38:01 +0200 Subject: [PATCH] CLI: Add color flag to `verdi` command (#6434) Adds the flag --color/--no-color to enforce color or no color for the output of the verdi commands. Implement support for NO_COLOR and FORCE_COLOR as specified in Python 3.13 for color commands. Implements feature request #4955: Color process states in output of `verdi process list`. --- docs/source/reference/command_line.rst | 1 + src/aiida/cmdline/commands/cmd_verdi.py | 1 + src/aiida/cmdline/groups/verdi.py | 13 +++- src/aiida/cmdline/params/options/__init__.py | 1 + src/aiida/cmdline/params/options/main.py | 28 +++++++- src/aiida/cmdline/utils/log.py | 5 +- src/aiida/common/style.py | 68 ++++++++++++++++++++ src/aiida/tools/query/formatting.py | 32 ++++++--- tests/cmdline/commands/test_process.py | 12 ++++ tests/cmdline/commands/test_verdi.py | 19 ++++++ tests/conftest.py | 1 + 11 files changed, 169 insertions(+), 12 deletions(-) create mode 100644 src/aiida/common/style.py diff --git a/docs/source/reference/command_line.rst b/docs/source/reference/command_line.rst index 0953d027f7..36fe76bf4e 100644 --- a/docs/source/reference/command_line.rst +++ b/docs/source/reference/command_line.rst @@ -176,6 +176,7 @@ Below is a list with all available subcommands. Options: -v, --verbosity [notset|debug|info|report|warning|error|critical] Set the verbosity of the output. + --color / --no-color Set if the output should be colorized. --help Show this message and exit. diff --git a/src/aiida/cmdline/commands/cmd_verdi.py b/src/aiida/cmdline/commands/cmd_verdi.py index c67b4b017a..ab33fd6182 100644 --- a/src/aiida/cmdline/commands/cmd_verdi.py +++ b/src/aiida/cmdline/commands/cmd_verdi.py @@ -20,6 +20,7 @@ @click.group(cls=VerdiCommandGroup, context_settings={'help_option_names': ['--help', '-h']}) @options.PROFILE(type=types.ProfileParamType(load_profile=True), expose_value=False) @options.VERBOSITY() +@options.COLOR() @click.version_option(__version__, package_name='aiida_core', message='AiiDA version %(version)s') def verdi(): """The command line interface of AiiDA.""" diff --git a/src/aiida/cmdline/groups/verdi.py b/src/aiida/cmdline/groups/verdi.py index e156b8413d..ccd5551fbd 100644 --- a/src/aiida/cmdline/groups/verdi.py +++ b/src/aiida/cmdline/groups/verdi.py @@ -136,6 +136,15 @@ def add_verbosity_option(cmd: click.Command) -> click.Command: return cmd + @staticmethod + def add_color_option(cmd: click.Command) -> click.Command: + """Apply the ``color`` option to the command, which is common to all ``verdi`` commands.""" + # Only apply the option if it hasn't been already added in a previous call. + if 'color' not in [param.name for param in cmd.params]: + cmd = options.COLOR()(cmd) + + return cmd + def fail_with_suggestions(self, ctx: click.Context, cmd_name: str) -> None: """Fail the command while trying to suggest commands to resemble the requested ``cmd_name``.""" # We might get better results with the Levenshtein distance or more advanced methods implemented in FuzzyWuzzy @@ -171,7 +180,9 @@ def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None cmd = super().get_command(ctx, cmd_name) if cmd is not None: - return self.add_verbosity_option(cmd) + cmd = self.add_verbosity_option(cmd) + cmd = self.add_color_option(cmd) + return cmd # If this command is called during tab-completion, we do not want to print an error message if the command can't # be found, but instead we want to simply return here. However, in a normal command execution, we do want to diff --git a/src/aiida/cmdline/params/options/__init__.py b/src/aiida/cmdline/params/options/__init__.py index 065efe4223..e4d623cb4e 100644 --- a/src/aiida/cmdline/params/options/__init__.py +++ b/src/aiida/cmdline/params/options/__init__.py @@ -103,6 +103,7 @@ 'USER_INSTITUTION', 'USER_LAST_NAME', 'VERBOSITY', + 'COLOR', 'VISUALIZATION_FORMAT', 'WAIT', 'WITH_ELEMENTS', diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index 85b3090ad5..45aed7cc66 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -107,6 +107,7 @@ 'USER_INSTITUTION', 'USER_LAST_NAME', 'VERBOSITY', + 'COLOR', 'VISUALIZATION_FORMAT', 'WAIT', 'WITH_ELEMENTS', @@ -175,7 +176,7 @@ def decorator(command): return decorator -def set_log_level(_ctx, _param, value): +def set_log_level(_ctx, _param, value) -> str: """Configure the logging for the CLI command being executed. Note that we cannot use the most obvious approach of directly setting the level on the various loggers. The reason @@ -226,6 +227,31 @@ def set_log_level(_ctx, _param, value): help='Set the verbosity of the output.', ) + +def set_color_option(ctx: click.Context, _, value: bool | None) -> bool: + """Sets the coloring for the CLI command outputs from given color option and returns if. + + :param ctx: The :class:`click.Command` that gives further information how the command was invoked. + :param value: The color option value given over the CLI. + """ + + from aiida.common.style import ColorConfig # We skip this when we are in a tab-completion context. + + if value is None and ctx.resilient_parsing: + return None + + ColorConfig.set_color(value) + return ColorConfig.get_color() + + +COLOR = OverridableOption( + '--color/--no-color', + default=None, + callback=set_color_option, + expose_value=False, # Ensures that the option is not actually passed to the command, because it doesn't need it + help='Set if the output should be colorized.', +) + PROFILE = OverridableOption( '-p', '--profile', diff --git a/src/aiida/cmdline/utils/log.py b/src/aiida/cmdline/utils/log.py index 46f8511ebc..ad79d0b1f5 100644 --- a/src/aiida/cmdline/utils/log.py +++ b/src/aiida/cmdline/utils/log.py @@ -4,6 +4,8 @@ import click +from aiida.common.style import ColorConfig + from .echo import COLORS @@ -35,7 +37,7 @@ def emit(self, record): try: msg = self.format(record) - click.echo(msg, err=err, nl=nl) + click.echo(msg, err=err, nl=nl, color=ColorConfig.get_color()) except Exception: self.handleError(record) @@ -59,5 +61,4 @@ def format(self, record): if prefix: return f'{click.style(record.levelname.capitalize(), fg=fg, bold=True)}: {formatted}' - return formatted diff --git a/src/aiida/common/style.py b/src/aiida/common/style.py new file mode 100644 index 0000000000..f605a32cdc --- /dev/null +++ b/src/aiida/common/style.py @@ -0,0 +1,68 @@ +########################################################################### +# Copyright (c), The AiiDA team. All rights reserved. # +# This file is part of the AiiDA code. # +# # +# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # +# For further information on the license, see the LICENSE.txt file # +# For further information please visit http://www.aiida.net # +########################################################################### +"""Utility functions to operate on datetime objects.""" + +import os +from typing import Optional + + +# Defines the styling for the process states +class ProcessStateStyle: + COLOR_CREATED_RUNNING = 'blue' + COLOR_WAITING = 'yellow' + COLOR_FINISHED = 'green' + COLOR_KILLED_EXPECTED = 'red' + + SYMBOL_EXPECTED = '\u2a2f' + SYMBOL_KILLED = '\u2620' + SYMBOL_CREATED_FINISHED = '\u23f9' + SYMBOL_RUNNING_WAITING = '\u23f5' + SYMBOL_RUNNING_WAITING_PAUSED = '\u23f8' + + +class ColorConfig: + """Controls the color styling option for aiida command outputs.""" + + _COLOR: bool | None = None + + @staticmethod + def get_color() -> bool | None: + """ + Returns the color value. If return value is None, the color value should be determined by caller. + """ + return ColorConfig._COLOR + + @staticmethod + def set_color(cli_color_option: Optional[bool] = None): + """ + Sets the color value that is determined from the CLI option or, if not + given, by the environment variables `FORCE_COLOR` and `NO_COLOR`. If also + no environment variable is given it set to `None` which signifies that + the caller of :meth:`~aiida.common.style.get_color` should determine if + output should allow colors. + + The logic for `FORCE_COLOR` and `NO_COLOR` follows the Python 3.13 implementation + See https://docs.python.org/3.13/using/cmdline.html#using-on-controlling-color + + :param cli_color_option: The option given over the CLI. + """ + ColorConfig._COLOR = None + + if cli_color_option is not None: + ColorConfig._COLOR = cli_color_option + else: + # Determines color for the terminal output depending on NO_COLOR and FORCE_COLOR + # environment variables following the Python implementation. + # See https://docs.python.org/3.13/using/cmdline.html#using-on-controlling-color + if os.getenv('TERM') == 'dump': + ColorConfig._COLOR = False + if 'FORCE_COLOR' in os.environ: + ColorConfig._COLOR = True + if 'NO_COLOR' in os.environ: + ColorConfig._COLOR = False diff --git a/src/aiida/tools/query/formatting.py b/src/aiida/tools/query/formatting.py index 5df70d34d0..153e94f32f 100644 --- a/src/aiida/tools/query/formatting.py +++ b/src/aiida/tools/query/formatting.py @@ -12,7 +12,10 @@ from datetime import datetime +import click + from aiida.common import timezone +from aiida.common.style import ColorConfig, ProcessStateStyle from aiida.common.utils import str_timedelta @@ -34,24 +37,37 @@ def format_state(process_state: str, paused: bool | None = None, exit_status: in :return: String representation of the process' state. """ if process_state in ['excepted']: - symbol = '\u2a2f' + symbol = ProcessStateStyle.SYMBOL_EXPECTED elif process_state in ['killed']: - symbol = '\u2620' + symbol = ProcessStateStyle.SYMBOL_KILLED elif process_state in ['created', 'finished']: - symbol = '\u23f9' + symbol = ProcessStateStyle.SYMBOL_CREATED_FINISHED elif process_state in ['running', 'waiting']: if paused is True: - symbol = '\u23f8' + symbol = ProcessStateStyle.SYMBOL_RUNNING_WAITING_PAUSED else: - symbol = '\u23f5' + symbol = ProcessStateStyle.SYMBOL_RUNNING_WAITING else: # Unknown process state, use invisible separator symbol = '\u00b7' # middle dot + output = f'{symbol} {format_process_state(process_state)}' if process_state == 'finished' and exit_status is not None: - return f'{symbol} {format_process_state(process_state)} [{exit_status}]' - - return f'{symbol} {format_process_state(process_state)}' + output += f' [{exit_status}]' + if ColorConfig.get_color(): + if process_state in ['created', 'running']: + color = ProcessStateStyle.COLOR_CREATED_RUNNING + elif process_state in ['waiting']: + color = ProcessStateStyle.COLOR_WAITING + elif process_state in ['finished']: + color = ProcessStateStyle.COLOR_FINISHED + elif process_state in ['killed', 'excepted']: + color = ProcessStateStyle.COLOR_KILLED_EXPECTED + else: + color = None + return click.style(output, color) + else: + return output def format_process_state(process_state: str | None) -> str: diff --git a/tests/cmdline/commands/test_process.py b/tests/cmdline/commands/test_process.py index fae9957f80..2f44bc662b 100644 --- a/tests/cmdline/commands/test_process.py +++ b/tests/cmdline/commands/test_process.py @@ -14,12 +14,14 @@ import typing as t import uuid +import click import pytest from aiida import get_profile from aiida.cmdline.commands import cmd_process from aiida.cmdline.utils.echo import ExitCode from aiida.common.links import LinkType from aiida.common.log import LOG_LEVEL_REPORT +from aiida.common.style import ProcessStateStyle from aiida.engine import Process, ProcessState from aiida.engine.processes import control as process_control from aiida.orm import CalcJobNode, Group, WorkChainNode, WorkflowNode, WorkFunctionNode @@ -183,6 +185,16 @@ def test_list(self, run_cli_command): result = run_cli_command(cmd_process.process_list, ['-r', '-X', flag, 'exit_message']) assert Process.exit_codes.ERROR_UNSPECIFIED.message in result.output + # check the color option works properly + colored_created = click.style( + f'{ProcessStateStyle.SYMBOL_CREATED_FINISHED} Created', ProcessStateStyle.COLOR_CREATED_RUNNING + ) + result = run_cli_command(cmd_process.process_list, ['--color']) + assert colored_created in result.output + + result = run_cli_command(cmd_process.process_list, ['--no-color']) + assert colored_created not in result.output + def test_process_show(self, run_cli_command): """Test verdi process show""" workchain_one = WorkChainNode() diff --git a/tests/cmdline/commands/test_verdi.py b/tests/cmdline/commands/test_verdi.py index 559d975515..9bc8c82251 100644 --- a/tests/cmdline/commands/test_verdi.py +++ b/tests/cmdline/commands/test_verdi.py @@ -65,3 +65,22 @@ def recursively_check_leaf_commands(ctx, command, leaf_commands): leaf_commands = [] ctx = click.Context(cmd_verdi.verdi) recursively_check_leaf_commands(ctx, cmd_verdi.verdi, leaf_commands) + + +def test_color_options(): + """Recursively find all leaf commands of ``verdi`` and ensure they have the ``--color`` option.""" + + def recursively_check_leaf_commands(ctx, command, leaf_commands): + """Recursively return the leaf commands of the given command.""" + try: + for subcommand in command.commands: + # We need to fetch the subcommand through the ``get_command``, because that is what the ``verdi`` + # command does when a subcommand is invoked on the command line. + recursively_check_leaf_commands(ctx, command.get_command(ctx, subcommand), leaf_commands) + except AttributeError: + # There are not subcommands so this is a leaf command, verify it has the color option + assert 'color' in [p.name for p in command.params], f'`{command.name} does not have color option' + + leaf_commands = [] + ctx = click.Context(cmd_verdi.verdi) + recursively_check_leaf_commands(ctx, cmd_verdi.verdi, leaf_commands) diff --git a/tests/conftest.py b/tests/conftest.py index 19d282548c..dacfcdfbb9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -724,6 +724,7 @@ def run_cli_command_runner(command, parameters, user_input, initialize_ctx_obj, # ``VerdiCommandGroup``, but when testing commands, the command is retrieved directly from the module which # circumvents this machinery. command = VerdiCommandGroup.add_verbosity_option(command) + command = VerdiCommandGroup.add_color_option(command) runner = CliRunner(mix_stderr=False) result = runner.invoke(command, parameters, input=user_input, obj=obj, **kwargs)