mujoco-mojo 2.3.6__tar.gz → 2.3.7__tar.gz

{mujoco_mojo-2.3.6 → mujoco_mojo-2.3.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mujoco-mojo
-Version: 2.3.6
+Version: 2.3.7
 Summary: A complete MJCF lifecycle and trial orchestration suite for MuJoCo, powered by Pydantic v2.
 Project-URL: Homepage, https://github.com/Hydrowelder/mujoco-mojo
 Project-URL: Documentation, https://hydrowelder.github.io/mujoco-mojo/

{mujoco_mojo-2.3.6 → mujoco_mojo-2.3.7}/src/mujoco_mojo/utils/layers/cli.py

@@ -4,14 +4,18 @@ import ast
 import importlib
 import logging
 import sys
+from collections.abc import Callable
 from enum import StrEnum
 from importlib.metadata import version
 from pathlib import Path
-from typing import Annotated, Any, Literal
+from types import ModuleType
+from typing import Annotated, Any, Literal, overload
 
 import typer
+from rich.align import Align
 from rich.console import Console
 from rich.panel import Panel
+from rich.text import Text
 
 # get logger is not called at the top of this module since it MUST be called after setup_logger is run
 # but since setup_logger doesnt know its verbosity until runtime get_logger needs to be called AS NEEDED
@@ -45,22 +49,37 @@ console = Console()
 VERSION = version("mujoco-mojo")
 
 
+# "MUJOCO" and "MOJO" in the ANSI Shadow figlet font
+_LOGO_LINES = [
+    "███╗   ███╗██╗   ██╗     ██╗ ██████╗  ██████╗ ██████╗       ███╗   ███╗ ██████╗      ██╗ ██████╗ ",
+    "████╗ ████║██║   ██║     ██║██╔═══██╗██╔════╝██╔═══██╗      ████╗ ████║██╔═══██╗     ██║██╔═══██╗",
+    "██╔████╔██║██║   ██║     ██║██║   ██║██║     ██║   ██║      ██╔████╔██║██║   ██║     ██║██║   ██║",
+    "██║╚██╔╝██║██║   ██║██   ██║██║   ██║██║     ██║   ██║      ██║╚██╔╝██║██║   ██║██   ██║██║   ██║",
+    "██║ ╚═╝ ██║╚██████╔╝╚█████╔╝╚██████╔╝╚██████╗╚██████╔╝      ██║ ╚═╝ ██║╚██████╔╝╚█████╔╝╚██████╔╝",
+    "╚═╝     ╚═╝ ╚═════╝  ╚════╝  ╚═════╝  ╚═════╝ ╚═════╝       ╚═╝     ╚═╝ ╚═════╝  ╚════╝  ╚═════╝ ",
+]
+
+_SHADES = [
+    "#67e8f9",
+    "#22d3ee",
+    "#06b6d4",
+    "#0891b2",
+    "#0e7490",
+    "#155e75",
+]
+
+
 def print_logo():
-    logo = """
-M   M  U  U  J   OOO    CCCC   OOO      M   M   OOO   J   OOO
-MM MM  U  U  J  OO OO  CC     OO OO     MM MM  OO OO  J  OO OO
-MM MM  U  U  J  O   O  C      O   O     MM MM  O   O  J  O   O
-MM MM  U  U  J  O   O  C      O   O     MM MM  O   O  J  O   O
-M M M  U  U  J  OO OO  CC  C  OO OO     M M M  OO OO  J  OO OO
-M M M  UUUU  J   OOO    CCCC   OOO      M M M   OOO   J   OOO
-             J                                        J
-           JJ                                       JJ
-"""
+    body = Text(justify="center")
+    for i, (line, shade) in enumerate(zip(_LOGO_LINES, _SHADES)):
+        body.append(line, style=f"bold {shade}")
+        if i != len(_LOGO_LINES) - 1:
+            body.append("\n")
+
     console.print(
         Panel(
-            logo,
+            Align.center(body),
             expand=False,
-            style="bold cyan",
             border_style="cyan",
             subtitle=f"[dim]v{VERSION}[/dim]",
         )
@@ -72,8 +91,22 @@ def _setup_cli_logging(verbose: int, quiet: int) -> logging.Logger:
     return setup_logger(level=level)
 
 
-def _load_func(path: str) -> Any:
-    """Helper to dynamically import user logic, supporting classes and methods."""
+@overload
+def _load_func(
+    path: str, _return_module: Literal[False] = False
+) -> Callable[..., Any]: ...
+@overload
+def _load_func(
+    path: str, _return_module: Literal[True]
+) -> tuple[Callable[..., Any], ModuleType]: ...
+def _load_func(
+    path: str, _return_module: bool = False
+) -> Callable[..., Any] | tuple[Callable[..., Any], ModuleType]:
+    """
+    Helper to dynamically import user logic, supporting classes and methods.
+
+    If `_return_module` is set, returns a `(obj, module)` tuple, where `module` is the importable module the dotted path resolved against (e.g. a package `__init__` re-exporting a class), as opposed to the module the attribute is actually defined in. Reloading that module is what's needed to pick up re-exported names after an edit.
+    """
     if "." not in sys.path:
         sys.path.insert(0, ".")
 
@@ -104,7 +137,6 @@ def _load_func(path: str) -> Any:
         obj = mod
         for part in attr_parts:
             obj = getattr(obj, part)
-        return obj
     except AttributeError as e:
         logger = get_logger(__name__)
         console.print(
@@ -113,6 +145,22 @@ def _load_func(path: str) -> Any:
         logger.exception(f"`{path}`not found: {e}", extra={"file_only": True})
         raise typer.Exit(code=1)
 
+    if not callable(obj):
+        logger = get_logger(__name__)
+        console.print(
+            f"[bold red]Error:[/bold red] [bold green]`{path}`[/bold green] resolved to a "
+            f"[bold]{type(obj).__name__}[/bold], which is not callable."
+        )
+        logger.error(
+            f"`{path}` resolved to a non-callable {type(obj).__name__}",
+            extra={"file_only": True},
+        )
+        raise typer.Exit(code=1)
+
+    if _return_module:
+        return obj, mod
+    return obj
+
 
 def _smart_parse(value: str) -> Any:
     """Tries to convert CLI strings to Python types (int, float, bool)."""

{mujoco_mojo-2.3.6 → mujoco_mojo-2.3.7}/src/mujoco_mojo/utils/layers/reloaded.py

@@ -23,7 +23,6 @@ from mujoco_mojo.stochas import DesignValueDict, DistributionDict, NamedValueDic
 from mujoco_mojo.utils.defaults import DEFAULT_WORKDIR
 from mujoco_mojo.utils.log import get_logger
 from mujoco_mojo.utils.runner import MojoGenerator, MojoRunner, MojoRuntime
-from mujoco_mojo.utils.utils import write_dojo_script
 from mujoco_mojo.utils.statusing import (
     JOB_STATUS_FNAME,
     TRIAL_STATUS_FNAME,
@@ -34,6 +33,7 @@ from mujoco_mojo.utils.statusing import (
     Step,
     TrialStatus,
 )
+from mujoco_mojo.utils.utils import write_dojo_script
 from mujoco_mojo.visualization import ArrowConfig, LineConfig
 
 from .cli import UserInterface
@@ -89,6 +89,44 @@ def is_dark_mode() -> bool:
     return False
 
 
+def _matrix_rain(duration: float = 1.5, fps: int = 20) -> None:
+    """Plays a brief Matrix-style digital rain animation, e.g. on exit."""
+    import random
+
+    from rich.live import Live
+    from rich.text import Text
+
+    width, height = console.size
+    if width <= 0 or height <= 0:
+        return
+
+    drops = [random.randint(-height, 0) for _ in range(width)]
+    trail_length = 8
+
+    with Live(console=console, refresh_per_second=fps, screen=True) as live:
+        for _ in range(int(duration * fps)):
+            frame = Text()
+            for row in range(height):
+                for col in range(width):
+                    head = drops[col]
+                    if row == head:
+                        frame.append(random.choice("01"), style="bold white")
+                    elif 0 <= head - row <= trail_length:
+                        frame.append(random.choice("01"), style="green")
+                    else:
+                        frame.append(" ")
+                if row < height - 1:
+                    frame.append("\n")
+
+            for col in range(width):
+                drops[col] += 1
+                if drops[col] - height > random.randint(0, height):
+                    drops[col] = random.randint(-height, 0)
+
+            live.update(frame)
+            time.sleep(1 / fps)
+
+
 def recursive_reload(module, project_root: Path, visited: set | None = None):
     """Recursively reloads modules, but only if they live inside project_root."""
     import importlib
@@ -271,9 +309,9 @@ class MojoReloaded:
         def reload_with_check(path_str: str, label: str):
             """Internal helper to validate file existance before reloading."""
             try:
-                # attempt to find
-                func = _load_func(path_str)
-                mod = sys.modules.get(func.__module__)
+                # attempt to find the module the dotted path resolves against, which may be a
+                # package `__init__` re-exporting a class/function defined in a submodule
+                _, mod = _load_func(path_str, _return_module=True)
 
                 if mod and hasattr(mod, "__file__") and mod.__file__:
                     source_file = Path(mod.__file__).resolve()
@@ -289,6 +327,8 @@ class MojoReloaded:
                             f"Unable to find source file for {path_str}. Did you rename or move it?"
                         )
 
+                    # reloading the entry module re-resolves any re-exported names (e.g.
+                    # `from .robot import RobotGenerator`) after their source modules reload
                     recursive_reload(module=mod, project_root=project_root)
 
                 # reload to get the reference
@@ -478,6 +518,7 @@ class MojoReloaded:
                 self.run_viser(initial_state)
             case _:
                 console.print(f"Invalid viewer option selected {self.ui}")
+        _matrix_rain()
         console.print("\n[bold yellow]Exiting MuJoCo Mojo Reloaded[/bold yellow]")
 
     def _print_help(self):
@@ -523,15 +564,18 @@ class MojoReloaded:
                     "Install [bold]mujoco-mojo[reloaded][/bold] to use --watch."
                 )
                 return None
+            # silence watchfiles' own "N changes detected" log; we print our own below
+            logging.getLogger("watchfiles").setLevel(logging.WARNING)
+
             evt = threading.Event()
             watch_path = Path.cwd()
 
             def _run():
                 try:
-                    for _ in wf_watch(
+                    for changes in wf_watch(
                         watch_path, watch_filter=PythonFilter(), stop_event=evt
                     ):
-                        event_queue.put(("reload", ""))
+                        event_queue.put(("reload", str(len(changes))))
                 except Exception:
                     pass
 
@@ -566,9 +610,7 @@ class MojoReloaded:
         def print_prompt():
             watch_indicator = " [bold blue]W[/bold blue]" if self.watch else ""
             console.print(
-                f"[bold green]Awaiting command[/bold green]"
-                f"[dim](last: {self._last_command} | seed: {self.seed} | trial: {self.trial_num})[/dim]"
-                f"{watch_indicator}[white] > [/white]",
+                f"[bold green]Awaiting command[/bold green] [dim](last: {self._last_command} | seed: {self.seed} | trial: {self.trial_num})[/dim] {watch_indicator}[white] > [/white]",
                 end="",
             )
 
@@ -583,8 +625,10 @@ class MojoReloaded:
                 break
 
             if event_type == "reload":
+                count = int(data) if data.isdigit() else 0
+                plural = "s" if count != 1 else ""
                 console.print(
-                    "\n[dim]File change detected, reloading with last command...[/dim]"
+                    f"\n[dim]{count} change{plural} detected, reloading with last command...[/dim]"
                 )
                 use_runtime = self._last_command != "gen"
                 try:
@@ -682,7 +726,7 @@ class MojoReloaded:
 
             try:
                 start = time.time()
-                console.print("[bold]Processing...[/bold]")
+                console.print("[dim]Processing...[/dim]")
                 new_state = self.generate_construct(
                     use_runtime=use_runtime, on_reload_callback=on_reload_callback
                 )