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

@@ -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.
 
 

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

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

src/aiida/cmdline/params/options/__init__.py

@@ -103,6 +103,7 @@
     'USER_INSTITUTION',
     'USER_LAST_NAME',
     'VERBOSITY',
+    'COLOR',
     'VISUALIZATION_FORMAT',
     'WAIT',
     'WITH_ELEMENTS',

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

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

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

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:

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

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)

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)