diff --git a/easybuild/base/testing.py b/easybuild/base/testing.py index 58c60d21df..b7de7033cf 100644 --- a/easybuild/base/testing.py +++ b/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 diff --git a/easybuild/tools/build_log.py b/easybuild/tools/build_log.py index e6a5103c27..9f1df4e8a6 100644 --- a/easybuild/tools/build_log.py +++ b/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="") diff --git a/easybuild/tools/config.py b/easybuild/tools/config.py index d6769442d9..2e42106b74 100644 --- a/easybuild/tools/config.py +++ b/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', ], diff --git a/easybuild/tools/entrypoints.py b/easybuild/tools/entrypoints.py index 1f7702d295..315b2a2b4c 100644 --- a/easybuild/tools/entrypoints.py +++ b/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__() diff --git a/easybuild/tools/options.py b/easybuild/tools/options.py index 96bb8a484a..796074b7fd 100644 --- a/easybuild/tools/options.py +++ b/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() diff --git a/easybuild/tools/output.py b/easybuild/tools/output.py index d1cf1810be..2dd78dd24f 100644 --- a/easybuild/tools/output.py +++ b/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)|(FAILED)|(FAIL))', + r'(?PWARNING)', + r'(?P(COMPLETED)|(SUCCESS)|(PASSED)|(PASS)|(OK))', + fr'(?P{EB_MSG_PREFIX} )', + r'(?P >> )', + r'.* (?P(building and installing|installing extension|installing bundle component).*)\.\.\.', + r'(?P\(took .*\))', + r'\[SUCCESS\] (?P\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. diff --git a/setup.py b/setup.py index 2f1a798a28..6ec2ef57ef 100644 --- a/setup.py +++ b/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/*')), diff --git a/test/framework/entrypoints.py b/test/framework/entrypoints.py index f052a6ee47..bf8b77bca9 100644 --- a/test/framework/entrypoints.py +++ b/test/framework/entrypoints.py @@ -37,10 +37,13 @@ from unittest import TextTestRunner import easybuild.tools.options as eboptions +import easybuild.tools.output as eboutput +from easybuild.tools.config import update_build_option from easybuild.tools.build_log import EasyBuildError from easybuild.tools.docs import list_easyblocks, list_toolchains from easybuild.tools.entrypoints import ( - HAVE_ENTRY_POINTS, EntrypointHook, EntrypointEasyblock, EntrypointToolchain, EasybuildEntrypoint + HAVE_ENTRY_POINTS, EntrypointHook, EntrypointEasyblock, EntrypointToolchain, EasybuildEntrypoint, + EntrypointRichTheme, EntrypointRichHighlighter ) from easybuild.tools.filetools import write_file from easybuild.tools.hooks import run_hook, START, CONFIGURE_STEP @@ -57,10 +60,24 @@ MOCK_HOOK_EP_NAME = "mock_hook" MOCK_EASYBLOCK_EP_NAME = "mock_easyblock" MOCK_TOOLCHAIN_EP_NAME = "mock_toolchain" +MOCK_THEME_EP_NAME = "mock_theme" +MOCK_HIGHLIGHTER_EP_NAME = "mock_highlighter" MOCK_HOOK = "hello_world_12412412" MOCK_EASYBLOCK = "TestEasyBlock_1212461" MOCK_TOOLCHAIN = "MockTc_352124671346" +MOCK_THEME = "MockTheme_124124" +MOCK_HIGHLIGHTER = "MockHighlighter_124124" + +MOCK_HIGHLIGHTER_TEXT = "TEST_TEXT" + +EXPECTED = { + EntrypointHook: MOCK_HOOK_EP_NAME, + EntrypointEasyblock: MOCK_EASYBLOCK_EP_NAME, + EntrypointToolchain: MOCK_TOOLCHAIN_EP_NAME, + EntrypointRichTheme: MOCK_THEME_EP_NAME, + EntrypointRichHighlighter: MOCK_HIGHLIGHTER_EP_NAME, +} MOCK_EP_FILE = f""" @@ -115,6 +132,21 @@ class {MOCK_TOOLCHAIN}(MockCompiler): class {MOCK_TOOLCHAIN}_invalid(MockCompiler): pass + +########################################################################## +from easybuild.tools.entrypoints import EntrypointRichTheme, EntrypointRichHighlighter + +@EntrypointRichTheme() +def {MOCK_THEME}(): + return {{ + 'easybuild.test': 'blue', + }} + +@EntrypointRichHighlighter() +def {MOCK_HIGHLIGHTER}(): + return [ + r'(?P{MOCK_HIGHLIGHTER_TEXT})', # Using `TEST` to distinguish this from other highlights + ] """ @@ -130,6 +162,12 @@ class {MOCK_TOOLCHAIN}_invalid(MockCompiler): [{EntrypointToolchain.group}] {MOCK_TOOLCHAIN_EP_NAME} = {{module}}:{MOCK_TOOLCHAIN} {{invalid_toolchain}} + +[{EntrypointRichTheme.group}] +{MOCK_THEME_EP_NAME} = {{module}}:{MOCK_THEME} + +[{EntrypointRichHighlighter.group}] +{MOCK_HIGHLIGHTER_EP_NAME} = {{module}}:{MOCK_HIGHLIGHTER} """ FORMAT_DCT = { @@ -288,29 +326,50 @@ class MOCK(Toolchain): pass decorator(MOCK) + def test_entrypoints_register_rich_theme(self): + """Test registering entry point rich themes with both valid and invalid theme names.""" + decorator = EntrypointRichTheme() + + with self.assertRaisesRegex(EasyBuildError, "has no module or name associated"): + decorator(123) + + with self.assertRaisesRegex(EasyBuildError, "did not return a dict"): + decorator(lambda: None) + + decorator(lambda: {}) + + def test_entrypoints_register_rich_highlighter(self): + """Test registering entry point rich highlighters with both valid and invalid highlighter names.""" + decorator = EntrypointRichHighlighter() + + with self.assertRaisesRegex(EasyBuildError, "has no module or name associated"): + decorator(123) + + with self.assertRaisesRegex(EasyBuildError, "does not return a list, got"): + decorator(lambda: None) + + with self.assertRaisesRegex(EasyBuildError, "did not return a list of strings"): + decorator(lambda: ['abc', 123]) + + decorator(lambda: ['a', 'b']) + def test_entrypoints_get_group(self): """Test retrieving entrypoints for a specific group.""" - expected = { - EntrypointHook: MOCK_HOOK_EP_NAME, - EntrypointEasyblock: MOCK_EASYBLOCK_EP_NAME, - EntrypointToolchain: MOCK_TOOLCHAIN_EP_NAME, - } - for ep_type in [EntrypointHook, EntrypointEasyblock, EntrypointToolchain]: + for ep_type in EXPECTED: group = ep_type.group epts = ep_type.retrieve_entrypoints() self.assertIsInstance(epts, set, f"Expected set for group {group}") self.assertEqual(len(epts), 0, f"Expected non-empty set for group {group}") init_config(build_options={'use_entrypoints': True}) - for ep_type in [EntrypointHook, EntrypointEasyblock, EntrypointToolchain]: + for ep_type, expt in EXPECTED.items(): group = ep_type.group epts = ep_type.retrieve_entrypoints() self.assertIsInstance(epts, set, f"Expected set for group {group}") self.assertGreater(len(epts), 0, f"Expected non-empty set for group {group}") loaded_names = [ep.name for ep in epts] - expt = expected[ep_type] self.assertIn(expt, loaded_names, f"Expected entry point {expt} in group {group}") def test_entrypoints_exclude_invalid(self): @@ -383,15 +442,18 @@ def test_entrypoints_show_config(self): args = ['--show-config'] stdout, stderr = self._run_mock_eb(args, strip=True) - for name in ['Hooks', 'Easyblocks', 'Toolchains']: - pattern = f"{name} from entrypoints (" + # for name in ['Hooks', 'Easyblocks', 'Toolchains']: + for ep_type in EXPECTED: + name = ep_type.desc + pattern = f"{name}s from entrypoints (" self.assertIn(pattern, stdout, f"Expected {name} in configuration output") args = ['--show-full-config'] stdout, stderr = self._run_mock_eb(args, strip=True) - for name in ['Hooks', 'Easyblocks', 'Toolchains']: - pattern = f"{name} from entrypoints (" + for ep_type in EXPECTED: + name = ep_type.desc + pattern = f"{name}s from entrypoints (" self.assertIn(pattern, stdout, f"Expected {name} in configuration output") def test_entrypoints_register_invalid_hook(self): @@ -477,6 +539,69 @@ def func_configure_post(): for key, val in flags.items(): self.assertEqual(val, key == 'post_cfg', "Should only run post-configure hooks") + def test_entrypoints_use_rich_theme(self): + """Test registering entry point rich themes with both valid and invalid theme names.""" + from easybuild.tools.build_log import print_msg + + update_build_option('output_style', 'rich') + + def run_check(msg, theme=None, hl=None, with_ep=False): + reload(eboutput) + update_build_option('output_theme', theme or eboutput.DEFAULT_THEME_NAME) + update_build_option('output_highlights', hl or eboutput.DEFAULT_HIGHLIGHTS_NAME) + update_build_option('use_entrypoints', with_ep) + + print_msg(msg) + + # Test that without entry points enabled, the custom theme and highlighter are not applied to the output + msg = "This is a test message" + with self.mocked_stdout_stderr(force_tty=True) as (stdout, _): + run_check(msg) + stdout_txt = stdout.getvalue().strip() + self.assertTrue(stdout_txt.endswith(f"\x1b[0m{msg}")) + + # Test that without entry points enabled, the custom theme and highlighter ONLY are applied + # The custom theme defined here should NOT be applied hece the TEST_TEXT should not be colored + msg = MOCK_HIGHLIGHTER_TEXT + with self.mocked_stdout_stderr(force_tty=True) as (stdout, _): + run_check(msg) + stdout_txt = stdout.getvalue().strip() + self.assertFalse(stdout_txt.endswith(f"{msg}\x1b[0m")) + + # Test that specifying a theme without entry points enabled raises an error with the expected message + reload(eboutput) # Clear out the cached theme/highlighter to ensure the new ones are picked up + exp = f'Cannot use custom Rich theme \'{MOCK_THEME}\' without entry points support enabled' + with self.assertRaisesRegex(EasyBuildError, exp): + run_check(msg, theme=MOCK_THEME) + + # Test that specifying an highlighter without entry points enabled raises an error with the expected message + exp = f'Cannot use custom Rich highlighter \'{MOCK_HIGHLIGHTER}\' without entry points support enabled' + with self.assertRaisesRegex(EasyBuildError, exp): + run_check(msg, hl=MOCK_HIGHLIGHTER) + + # Test that specifying an invalid highlighter with entry points enabled raises an error with the expected + # message listing the available highlighters + highlighter = f'{MOCK_HIGHLIGHTER}_invalid' + exp = rf'Unknown specified Rich highlighter.*{highlighter}.* \(available: .*\)' + with self.assertRaisesRegex(EasyBuildError, exp): + run_check(msg, hl=highlighter, with_ep=True) + + # Test that specifying an invalid theme with entry points enabled raises an error with the expected message + # listing the available themes + theme = f'{MOCK_THEME}_invalid' + exp = rf'Unknown specified Rich theme.*{theme}.* \(available: .*\)' + with self.assertRaisesRegex(EasyBuildError, exp): + run_check(msg, theme=theme, with_ep=True) + + # Test that specifying a valid highlighter with entry points enabled does not raise an error and is applied + # to the output + highlighter = MOCK_HIGHLIGHTER + exp = rf'Unknown specified Rich highlighter.*{highlighter}.* \(available: .*\)' + with self.mocked_stdout_stderr(force_tty=True) as (stdout, _): + run_check(msg, theme=MOCK_THEME, hl=MOCK_HIGHLIGHTER, with_ep=True) + stdout_txt = stdout.getvalue().strip() + self.assertTrue(stdout_txt.endswith(f"{msg}\x1b[0m")) + def suite(loader=None): """ returns all the testcases in this module """