Allow setting the terminal effect in hydra_command

Signed-off-by: guillemdb <guillem.duran@inait.ai>

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

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", "")

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)