diff --git a/src/hydraclick/__init__.py b/src/hydraclick/__init__.py index 7b9d442..5893bbc 100644 --- a/src/hydraclick/__init__.py +++ b/src/hydraclick/__init__.py @@ -1,2 +1,2 @@ -from hydraclick.terminal_effects import set_terminal_effect, NO_TERMINAL_EFFECTS +from hydraclick.terminal_effects import set_terminal_effect from hydraclick.core import hydra_command diff --git a/src/hydraclick/core.py b/src/hydraclick/core.py index 5dcf6b8..68bff42 100644 --- a/src/hydraclick/core.py +++ b/src/hydraclick/core.py @@ -5,10 +5,11 @@ from typing import Callable, Any import hydra +import omegaconf from omegaconf import DictConfig, OmegaConf from unittest.mock import patch - +from hydraclick import set_terminal_effect from hydraclick.display_config import display_config from hydraclick.options import ( hydra_args_argument, @@ -25,6 +26,7 @@ config_name_option, shell_completion_option, ) +from hydraclick.terminal_effects import display_terminal_effect _logger = logging.getLogger(__name__) @@ -171,6 +173,7 @@ def command_api( print_config: bool = True, resolve: bool = True, use_flogging: bool = True, + terminal_effect: Callable | None = omegaconf.MISSING, **flogging_kwargs: Any, ) -> Callable: """Integrate Hydra's configuration management capabilities with a Click-based CLI. @@ -197,6 +200,8 @@ def command_api( function. Defaults to `True`. use_flogging (bool, optional): Whether to use the `flogging` library for structured \ logging. Defaults to `True`. + terminal_effect(Callable | None, optional): the terminal effect function to use when \ + rendering the command help. **flogging_kwargs (Any, optional): Additional keyword arguments to pass to the \ `flogging.setup` function. @@ -237,6 +242,10 @@ def my_function(config: DictConfig): configuration before it is passed to the main function. """ + if terminal_effect == omegaconf.MISSING: + terminal_effect = display_terminal_effect + if terminal_effect is not None: + set_terminal_effect(terminal_effect) config_path = get_default_dir() if config_path is None else str(config_path) if config_name is not None: config_name = str(config_name).replace(".yaml", "").replace(".yml", "") diff --git a/src/hydraclick/terminal_effects.py b/src/hydraclick/terminal_effects.py index 12efb7f..8271f38 100644 --- a/src/hydraclick/terminal_effects.py +++ b/src/hydraclick/terminal_effects.py @@ -9,23 +9,18 @@ def get_no_terminal_efects() -> bool: """Get the no terminal effects environment variable.""" - val = os.environ.get( - "OMEGACLICK_NO_TERMINAL_EFFECTS", os.environ.get("OMEGACLICK_NO_TERMINAL_EFFECTS") - ) + val = os.environ.get("OMEGACLICK_NO_TERMINAL_EFFECTS", os.environ.get("NO_TERMINAL_EFFECTS")) if val is None: return False return val.lower() in {"true", "1", "yes"} -NO_TERMINAL_EFFECTS = get_no_terminal_efects() - - def config_effect(effect): """Configure the terminal effect.""" from terminaltexteffects.utils.graphics import Color # noqa: PLC0415 - effect.effect_config.print_speed = 5 - effect.effect_config.print_head_return_speed = 3 + effect.effect_config.print_speed = 15 + effect.effect_config.print_head_return_speed = 5 effect.effect_config.final_gradient_stops = (Color("00ffae"), Color("00D1FF"), Color("FFFFFF")) return effect @@ -75,59 +70,57 @@ def display_terminal_effect(value, effect_cls=None): sys.stdout.flush() -_TERMINAL_EFFECT = display_terminal_effect - - -def set_terminal_effect(value: Callable | None = None): - """Set the terminal effect.""" - global _TERMINAL_EFFECT # noqa: PLW0603 - if value is None: - value = display_terminal_effect - _TERMINAL_EFFECT = value +def patch_parse_args(terminal_effect: Callable): + """Patch the click `get_help_option` function with a custom terminal effect.""" + def parse_args(self, ctx: Context, args: list[str]) -> list[str]: + """Display the help message when no arguments are provided.""" + if not args and self.no_args_is_help and not ctx.resilient_parsing: + terminal_effect(ctx.get_help()) + ctx.exit() -def parse_args(self, ctx: Context, args: list[str]) -> list[str]: - """Display the help message when no arguments are provided.""" - if not args and self.no_args_is_help and not ctx.resilient_parsing: - _TERMINAL_EFFECT(ctx.get_help()) - ctx.exit() + rest = super(click.core.MultiCommand, self).parse_args(ctx, args) + if self.chain: + ctx.protected_args = rest + ctx.args = [] + elif rest: + ctx.protected_args, ctx.args = rest[:1], rest[1:] + return ctx.args - rest = super(click.core.MultiCommand, self).parse_args(ctx, args) + click.core.MultiCommand.parse_args = parse_args - if self.chain: - ctx.protected_args = rest - ctx.args = [] - elif rest: - ctx.protected_args, ctx.args = rest[:1], rest[1:] - return ctx.args +def patch_get_help_option(terminal_effect: Callable): + """Patch the click `get_help_option` function with a custom terminal effect.""" + def get_help_option(self, ctx: Context) -> Optional["Option"]: + """Return the help option object.""" + from gettext import gettext # noqa: PLC0415 -def get_help_option(self, ctx: Context) -> Optional["Option"]: - """Return the help option object.""" - from gettext import gettext # noqa: PLC0415 + help_options = self.get_help_option_names(ctx) - help_options = self.get_help_option_names(ctx) + if not help_options or not self.add_help_option: + return None - if not help_options or not self.add_help_option: - return None + def show_help(ctx: Context, param: "click.Parameter", value: str) -> None: # noqa: ARG001 + if value and not ctx.resilient_parsing: + terminal_effect(ctx.get_help()) + ctx.exit() - def show_help(ctx: Context, param: "click.Parameter", value: str) -> None: # noqa: ARG001 - if value and not ctx.resilient_parsing: - display_terminal_effect(ctx.get_help()) - ctx.exit() + return Option( + help_options, + is_flag=True, + is_eager=True, + expose_value=False, + callback=show_help, + help=gettext("Show this message and exit."), + ) - return Option( - help_options, - is_flag=True, - is_eager=True, - expose_value=False, - callback=show_help, - help=gettext("Show this message and exit."), - ) + click.core.Command.get_help_option = get_help_option -if not NO_TERMINAL_EFFECTS: - set_terminal_effect(display_terminal_effect) - click.core.Command.get_help_option = get_help_option - click.core.MultiCommand.parse_args = parse_args +def set_terminal_effect(terminal_effect): + """Set the terminal effect animation to appear when displaying the help.""" + if not get_no_terminal_efects(): + patch_parse_args(terminal_effect) + patch_get_help_option(terminal_effect)