Add customizable Themes and Highlighters for rich colors in print_msg

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,8 @@
 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'
+
 OUTPUT_STYLE_AUTO = 'auto'
 OUTPUT_STYLE_BASIC = 'basic'
 OUTPUT_STYLE_NO_COLOR = 'no_color'
@@ -437,6 +439,9 @@ def mk_full_default_path(name, prefix=DEFAULT_PREFIX):
     DEFAULT_PR_TARGET_ACCOUNT: [
         'pr_target_account',
     ],
+    DEFAULT_THEME_NAME: [
+        'output_theme',
+    ],
     GENERAL_CLASS: [
         'suffix_modules_path',
     ],

easybuild/tools/entrypoints.py

@@ -139,7 +139,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,6 +177,11 @@ 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'
+
+
 class EntrypointHook(EasybuildEntrypoint):
     """Class to represent a hook entrypoint."""
     group = 'easybuild.hooks'

easybuild/tools/options.py

@@ -555,6 +555,7 @@ 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'),
             '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)",

easybuild/tools/output.py

@@ -35,11 +35,14 @@
 from collections import OrderedDict
 import sys
 
-from easybuild.tools.build_log import EasyBuildError
-from easybuild.tools.config import OUTPUT_STYLE_RICH, build_option, get_output_style
+from easybuild.tools.build_log import EasyBuildError, EB_MSG_PREFIX
+from easybuild.tools.entrypoints import EntrypointRichTheme
+from easybuild.tools.config import OUTPUT_STYLE_RICH, build_option, get_output_style, DEFAULT_THEME_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 +78,35 @@
 
 _progress_bar_cache = {}
 
+DEFAULT_THEME_DCT = {
+    'easybuild.success': 'on green',
+    'easybuild.warning': 'on orange3',
+    'easybuild.error': 'on red',
+    'easybuild.prefix1': 'gray50',
+    'easybuild.prefix2': 'gray50',
+    'easybuild.installing': 'bold',
+    'easybuild.installed': 'bold',
+    'easybuild.timing': 'gray50',
+    'repr.path': 'bright_blue',
+    'repr.number': 'red',
+    'repr.ipv6': 'yellow',  # time
+    'repr.call': 'none',  # file(s)
+    'repr.ellipsis': 'gray50',
+    '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 +165,73 @@ def use_rich():
     return get_output_style() == OUTPUT_STYLE_RICH
 
 
+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)
+        if output_theme == DEFAULT_THEME_NAME:
+            theme_dct = DEFAULT_THEME_DCT
+        else:
+            if not use_entrypoints:
+                raise EasyBuildError(
+                    "Cannot use custom Rich theme '%s' without entry points support enabled", output_theme
+                )
+            for entrypoint in EntrypointRichTheme.retrieve_entrypoints():
+                if entrypoint.name == output_theme:
+                    theme_dct = entrypoint.load()
+                    break
+            else:
+                raise EasyBuildError("Unknown specified Rich theme '%s'", output_theme)
+        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:
+        class EasybuildHighlighter(RegexHighlighter):
+            """Highlighter for EasyBuild messages, to highlight ERROR, WARNING, SUCCESS and similar lines."""
+            highlights = DEFAULT_HIGHLIGHTS
+            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,11 @@ def find_rel_test():
         # utility scripts
         'easybuild/scripts/install_eb_dep.sh',
     ],
+    entry_points={
+        'easybuild.rich_theme': [
+            'default = easybuild.tools.output:DEFAULT_THEME_DCT',
+        ],
+    },
     data_files=[
         ('easybuild/scripts', glob.glob('easybuild/scripts/*')),
         ('etc', glob.glob('etc/*')),