Add customizable Themes and Highlighters for rich colors in print_msg

easybuild/base/testing.py

@@ -227,12 +227,12 @@ def mocked_stderr(self):
             self.mock_stderr(False)
 
     @contextmanager
-    def mocked_stdout_stderr(self, mock_stdout=True, mock_stderr=True):
+    def mocked_stdout_stderr(self, mock_stdout=True, mock_stderr=True, force_tty=False):
         """Context manager to mock stdout and stderr"""
         if mock_stdout:
-            self.mock_stdout(True)
+            self.mock_stdout(True, force_tty=force_tty)
         if mock_stderr:
-            self.mock_stderr(True)
+            self.mock_stderr(True, force_tty=force_tty)
         try:
             if mock_stdout and mock_stderr:
                 yield sys.stdout, sys.stderr

easybuild/tools/build_log.py

@@ -314,7 +314,7 @@ def print_msg(msg, *args, **kwargs):
     :param newline: end message with newline
     :param stderr: print to stderr rather than stdout
     """
-    from easybuild.tools.output import use_rich  # avoid circular import
+    from easybuild.tools.output import use_rich, get_rich_highlighter, get_rich_theme  # avoid circular import
     if args:
         msg = msg % args
 
@@ -339,7 +339,7 @@ def print_msg(msg, *args, **kwargs):
             from rich.markup import escape
             from rich.console import Console
 
-            console = Console()
+            console = Console(highlighter=get_rich_highlighter(), theme=get_rich_theme())
             with console.capture() as capture:
                 console.print(escape(msg), end="")
 

easybuild/tools/config.py

@@ -178,6 +178,9 @@
 LOCAL_VAR_NAMING_CHECK_WARN = WARN
 LOCAL_VAR_NAMING_CHECKS = [LOCAL_VAR_NAMING_CHECK_ERROR, LOCAL_VAR_NAMING_CHECK_LOG, LOCAL_VAR_NAMING_CHECK_WARN]
 
+DEFAULT_THEME_NAME = 'default_theme'
+DEFAULT_HIGHLIGHTS_NAME = 'default_highlights'
+
 OUTPUT_STYLE_AUTO = 'auto'
 OUTPUT_STYLE_BASIC = 'basic'
 OUTPUT_STYLE_NO_COLOR = 'no_color'
@@ -437,6 +440,12 @@ def mk_full_default_path(name, prefix=DEFAULT_PREFIX):
     DEFAULT_PR_TARGET_ACCOUNT: [
         'pr_target_account',
     ],
+    DEFAULT_THEME_NAME: [
+        'output_theme',
+    ],
+    DEFAULT_HIGHLIGHTS_NAME: [
+        'output_highlights',
+    ],
     GENERAL_CLASS: [
         'suffix_modules_path',
     ],

easybuild/tools/entrypoints.py

@@ -59,6 +59,7 @@ class EasybuildEntrypoint:
     group = None
     expected_type = None
     registered = {}
+    desc = None
 
     def __init__(self):
         if self.group is None:
@@ -139,7 +140,7 @@ def retrieve_entrypoints(cls) -> Set[EntryPoint]:
 
     @classmethod
     def load_entrypoints(cls):
-        """Load all the entrypoints in this group. This is needed for the modules contining the entrypoints to be
+        """Load all the entrypoints in this group. This is needed for the modules containing the entrypoints to be
         actually imported in order to process the function decorators that will register them in the
         `registered` dict."""
         for ep in cls.retrieve_entrypoints():
@@ -177,9 +178,39 @@ def validate(self):
             raise EasyBuildError("Entrypoint `%s` has no module or name associated", self.wrapped)
 
 
+class EntrypointRichTheme(EasybuildEntrypoint):
+    """Class to represent a rich theme entrypoint."""
+    group = 'easybuild.rich_theme'
+    desc = "Rich theme"
+
+    def validate(self):
+        super().validate()
+        if not callable(self.wrapped):
+            raise EasyBuildError("Rich theme entrypoint `%s` is not callable", self.wrapped)
+        res = self.wrapped()
+        if not isinstance(res, dict):
+            raise EasyBuildError("Rich theme entrypoint `%s` did not return a dict", self.wrapped)
+
+
+class EntrypointRichHighlighter(EasybuildEntrypoint):
+    """Class to represent a rich highlighter entrypoint."""
+    group = 'easybuild.rich_highlighter'
+    desc = "Rich highlighter"
+
+    def validate(self):
+        super().validate()
+        if not callable(self.wrapped):
+            raise EasyBuildError("Rich highlighter entrypoint `%s` is not callable", self.wrapped)
+        res = self.wrapped()
+        if not isinstance(res, list):
+            raise EasyBuildError("Rich highlighter entrypoint does not return a list, got %s", type(res))
+        if any(not isinstance(item, str) for item in res):
+            raise EasyBuildError("Rich highlighter entrypoint `%s` did not return a list of strings", self.wrapped)
+
 class EntrypointHook(EasybuildEntrypoint):
     """Class to represent a hook entrypoint."""
     group = 'easybuild.hooks'
+    desc = "Hook"
 
     def __init__(self, step, pre_step=False, post_step=False, priority=0):
         """Initialize the EntrypointHook."""
@@ -214,6 +245,7 @@ def validate(self):
 class EntrypointEasyblock(EasybuildEntrypoint):
     """Class to represent an easyblock entrypoint."""
     group = 'easybuild.easyblock'
+    desc = "Easyblock"
 
     def __init__(self):
         super().__init__()
@@ -225,6 +257,7 @@ def __init__(self):
 class EntrypointToolchain(EasybuildEntrypoint):
     """Class to represent a toolchain entrypoint."""
     group = 'easybuild.toolchain'
+    desc = "Toolchain"
 
     def __init__(self, prepend=False):
         super().__init__()

easybuild/tools/options.py

@@ -76,6 +76,7 @@
 from easybuild.tools.config import EBROOT_ENV_VAR_ACTIONS, ERROR, FORCE_DOWNLOAD_CHOICES, GENERAL_CLASS, IGNORE
 from easybuild.tools.config import JOB_DEPS_TYPE_ABORT_ON_ERROR, JOB_DEPS_TYPE_ALWAYS_RUN, LOADED_MODULES_ACTIONS
 from easybuild.tools.config import LOCAL_VAR_NAMING_CHECK_WARN, LOCAL_VAR_NAMING_CHECKS, MOD_SEARCH_PATH_HEADERS
+from easybuild.tools.config import DEFAULT_THEME_NAME, DEFAULT_HIGHLIGHTS_NAME
 from easybuild.tools.config import OUTPUT_STYLE_AUTO, OUTPUT_STYLES, WARN, build_option
 from easybuild.tools.config import get_pretend_installpath, init, init_build_options, mk_full_default_path
 from easybuild.tools.config import BuildOptions, ConfigurationVariables
@@ -110,7 +111,7 @@
 from easybuild.tools.systemtools import get_cpu_features, get_gpu_info, get_os_type, get_system_info
 from easybuild.tools.utilities import flatten
 from easybuild.tools.version import this_is_easybuild
-from easybuild.tools.entrypoints import EntrypointHook, EntrypointEasyblock, EntrypointToolchain
+from easybuild.tools.entrypoints import EasybuildEntrypoint
 
 
 try:
@@ -555,6 +556,10 @@ def override_options(self):
             'output-style': ("Control output style; auto implies using Rich if available to produce rich output, "
                              "with fallback to basic colored output",
                              'choice', 'store', OUTPUT_STYLE_AUTO, OUTPUT_STYLES),
+            'output-theme': ("Set output theme (when using Rich output style)", None, 'store', DEFAULT_THEME_NAME),
+            'output-highlights': (
+                "Set output highlights (when using Rich output style)", None, 'store', DEFAULT_HIGHLIGHTS_NAME
+            ),
             'parallel': ("Specify level of parallelism that should be used during build procedure, "
                          "(bypasses auto-detection of number of available cores; "
                          "actual value is determined by this value + 'max_parallel' easyconfig parameter)",
@@ -1685,11 +1690,8 @@ def det_location(opt, prefix=''):
         pretty_print_opts(opts_dict)
 
         if build_option('use_entrypoints', default=True):
-            for prefix, cls in [
-                ('Hook', EntrypointHook),
-                ('Easyblock', EntrypointEasyblock),
-                ('Toolchain', EntrypointToolchain),
-            ]:
+            for cls in EasybuildEntrypoint.__subclasses__():
+                prefix = str(cls.desc)
                 ept_list = cls.retrieve_entrypoints()
                 if ept_list:
                     print()

easybuild/tools/output.py

@@ -35,11 +35,15 @@
 from collections import OrderedDict
 import sys
 
-from easybuild.tools.build_log import EasyBuildError
+from easybuild.tools.build_log import EasyBuildError, EB_MSG_PREFIX
+from easybuild.tools.entrypoints import EntrypointRichTheme, EntrypointRichHighlighter
 from easybuild.tools.config import OUTPUT_STYLE_RICH, build_option, get_output_style
+from easybuild.tools.config import DEFAULT_THEME_NAME, DEFAULT_HIGHLIGHTS_NAME
 
 try:
     import rich.markup
+    from rich.theme import Theme
+    from rich.highlighter import Highlighter, RegexHighlighter, ReprHighlighter
     from rich.console import Console, Group
     from rich.live import Live
     from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
@@ -75,6 +79,35 @@
 
 _progress_bar_cache = {}
 
+DEFAULT_THEME_DCT = {
+    'easybuild.success': 'on green',
+    'easybuild.warning': 'on orange3',
+    'easybuild.error': 'on red',
+    'easybuild.prefix1': 'grey50',
+    'easybuild.prefix2': 'grey50',
+    'easybuild.installing': 'bold',
+    'easybuild.installed': 'bold',
+    'easybuild.timing': 'grey50',
+    'repr.path': 'bright_blue',
+    'repr.number': 'red',
+    'repr.ipv6': 'yellow',  # time
+    'repr.call': 'none',  # file(s)
+    'repr.ellipsis': 'grey50',
+    'repr.brace': 'none',
+}
+DEFAULT_HIGHLIGHTS = [
+    r'(?P<error>(ERROR)|(FAILED)|(FAIL))',
+    r'(?P<warning>WARNING)',
+    r'(?P<success>(COMPLETED)|(SUCCESS)|(PASSED)|(PASS)|(OK))',
+    fr'(?P<prefix1>{EB_MSG_PREFIX} )',
+    r'(?P<prefix2> >> )',
+    r'.* (?P<installing>(building and installing|installing extension|installing bundle component).*)\.\.\.',
+    r'(?P<timing>\(took .*\))',
+    r'\[SUCCESS\] (?P<installed>\S+)',
+]
+CACHED_THEME = None
+CACHED_HIGHLIGHTER = None
+
 
 def colorize(txt, color):
     """
@@ -133,6 +166,101 @@ def use_rich():
     return get_output_style() == OUTPUT_STYLE_RICH
 
 
+@EntrypointRichTheme()
+def default_theme():
+    """
+    Default Rich theme, used if no custom theme is specified or available.
+    """
+    return DEFAULT_THEME_DCT
+
+
+@EntrypointRichHighlighter()
+def default_highlights():
+    """
+    Default Rich highlighter, used if no custom highlighter is specified or available.
+    """
+    return DEFAULT_HIGHLIGHTS
+
+
+def get_rich_theme():
+    """
+    Get Rich theme to use for rich output.
+    """
+    global CACHED_THEME
+    if CACHED_THEME is not None:
+        return CACHED_THEME
+    if not use_rich():
+        class DummyTheme:
+            pass
+        res = DummyTheme()
+    else:
+        use_entrypoints = build_option('use_entrypoints', default=True)
+        output_theme = build_option('output_theme', default=DEFAULT_THEME_NAME)
+
+        entrypoints = EntrypointRichTheme.get_loaded_entrypoints(name=output_theme)
+        if not entrypoints:
+            if use_entrypoints:
+                available_themes = ', '.join([_.name for _ in EntrypointRichTheme.get_loaded_entrypoints()])
+                msg = f"Unknown specified Rich theme '{output_theme}' (available: {available_themes})"
+            else:
+                msg = f"Cannot use custom Rich theme '{output_theme}' without entry points support enabled"
+            raise EasyBuildError(msg)
+        theme_dct = entrypoints[0].wrapped()
+        res = Theme(theme_dct)
+
+    CACHED_THEME = res
+    return res
+
+
+def get_rich_highlighter():
+    """
+    Get Rich highlighter to use for rich output.
+    """
+    global CACHED_HIGHLIGHTER
+    if CACHED_HIGHLIGHTER is not None:
+        return CACHED_HIGHLIGHTER
+
+    if not use_rich():
+        class DummyHighlighter:
+            pass
+        res = DummyHighlighter()
+    else:
+        use_entrypoints = build_option('use_entrypoints', default=True)
+        output_hl = build_option('output_highlights', default=DEFAULT_HIGHLIGHTS_NAME)
+
+        entrypoints = EntrypointRichHighlighter.get_loaded_entrypoints(name=output_hl)
+        if not entrypoints:
+            if use_entrypoints:
+                available_hls = ', '.join([_.name for _ in EntrypointRichHighlighter.get_loaded_entrypoints()])
+                msg = f"Unknown specified Rich highlighter '{output_hl}' (available: {available_hls})"
+            else:
+                msg = f"Cannot use custom Rich highlighter '{output_hl}' without entry points support enabled"
+            raise EasyBuildError(msg)
+        highlights_dct = entrypoints[0].wrapped()
+
+        class EasybuildHighlighter(RegexHighlighter):
+            """Highlighter for EasyBuild messages, to highlight ERROR, WARNING, SUCCESS and similar lines."""
+            highlights = highlights_dct
+            base_style = "easybuild."
+
+        class CombinedHighlighter(Highlighter):
+            """Combined highlighter that applies both EasybuildHighlighter and ReprHighlighter."""
+            def __init__(self):
+                super().__init__()
+                self.easybuild_highlighter = EasybuildHighlighter()
+                self.repr_highlighter = ReprHighlighter()
+
+            def highlight(self, text):
+                self.repr_highlighter.highlight(text)
+                # easybuild_highlighter comes last to take priority
+                self.easybuild_highlighter.highlight(text)
+
+        res = CombinedHighlighter()
+
+    CACHED_HIGHLIGHTER = res
+    return res
+
+
 def show_progress_bars():
     """
     Return whether or not to show progress bars.

setup.py

@@ -119,6 +119,14 @@ def find_rel_test():
         # utility scripts
         'easybuild/scripts/install_eb_dep.sh',
     ],
+    entry_points={
+        'easybuild.rich_theme': [
+            'default = easybuild.tools.output:default_theme',
+        ],
+        'easybuild.rich_highlighter': [
+            'default = easybuild.tools.output:default_highlights',
+        ],
+    },
     data_files=[
         ('easybuild/scripts', glob.glob('easybuild/scripts/*')),
         ('etc', glob.glob('etc/*')),