relationalai 1.0.0a3__py3-none-any.whl → 1.0.0a5__py3-none-any.whl

v0/relationalai/tools/cli_controls.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 # Standard library imports
+import contextvars
 import io
 import itertools
 import os
@@ -12,7 +13,7 @@ import time
 import importlib
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Sequence, TextIO, cast
+from typing import Any, Callable, Dict, List, Optional, Sequence, TextIO, cast
 
 # Third-party imports
 import rich
@@ -36,6 +37,27 @@ from ..environments import (
     runtime_env,
 )
 
+# ---------------------------------------------
+# Global controls for nesting TaskProgress
+# ---------------------------------------------
+
+# Type alias for any progress type that supports nesting
+_ProgressType = Any  # Actually TaskProgress | NotebookTaskProgress, but defined before those classes
+
+_current_progress: contextvars.ContextVar[Optional[_ProgressType]] = contextvars.ContextVar(
+    'current_progress', default=None
+)
+
+
+def get_current_progress() -> Optional[_ProgressType]:
+    """Get the currently active TaskProgress, if any."""
+    return _current_progress.get()
+
+
+def _set_current_progress(progress: Optional[_ProgressType]) -> contextvars.Token:
+    """Set the current TaskProgress and return a token for restoration."""
+    return _current_progress.set(progress)
+
 #--------------------------------------------------
 # Constants
 #--------------------------------------------------
@@ -895,8 +917,14 @@ class _TaskStateMixin:
         self.add_sub_task(sub_status, "default")
         self.update_sub_task("default", sub_status)
 
+    def update_main_status_fn(self, fn: Callable[[], str]) -> None:
+        """Update the main status line using a callable function."""
+        self._description_fn = fn
+
     def update_main_status(self, message: str) -> None:
         """Update the main status line with custom information."""
+        if self._description_fn is not None:
+            self._description_fn = None
         if getattr(self, "description", "") != message:
             self.description = message
             self._after_task_update()
@@ -904,6 +932,8 @@ class _TaskStateMixin:
     def update_messages(self, updater: dict[str, str]) -> None:
         """Update both main message and sub-status if provided."""
         if "message" in updater:
+            if self._description_fn is not None:
+                self._description_fn = None
             self.description = updater["message"]
             self._after_task_update()
         if "sub_status" in updater:
@@ -998,7 +1028,7 @@ class TaskProgress(_TaskStateMixin):
 
     def __init__(
         self,
-        description: str = "",
+        description: str | Callable[[], str] = "",
         success_message: str = "",
         failure_message: str = "",
         leading_newline: bool = False,
@@ -1007,8 +1037,13 @@ class TaskProgress(_TaskStateMixin):
         hide_on_completion: bool = False,
         show_duration_summary: bool = True,
     ):
-        # Public configuration
-        self.description = description
+        # Public configuration - description can be a string or callable
+        if callable(description):
+            self._description_fn: Callable[[], str] | None = description
+            self.description = description()  # Initial value
+        else:
+            self._description_fn = None
+            self.description = description
         self.success_message = success_message
         self.failure_message = failure_message
         self.leading_newline = leading_newline
@@ -1046,19 +1081,29 @@ class TaskProgress(_TaskStateMixin):
 
         self._spinner_thread = None
 
+        # Nesting support
+        self._context_token: Optional[contextvars.Token] = None
+        self._parent_progress: Optional[TaskProgress] = None
+
     def _generate_task_id(self) -> str:
         """Generate a unique task ID."""
         task_id = f"task_{self._next_task_id}"
         self._next_task_id += 1
         return task_id
 
-    def _compute_state_hash(self) -> int:
+    def _get_description(self) -> str:
+        """Get the current description, calling the function if one was provided."""
+        if self._description_fn is not None:
+            return self._description_fn()
+        return self.description
+
+    def _compute_state_hash(self, description: str) -> int:
         """Compute a simple hash of the current state for caching."""
         # Use a simple hash based on key state variables
         state_parts = [
             str(self.main_completed),
             str(self.main_failed),
-            self.description,
+            description,
             str(self.spinner_index),
             str(len(self._tasks)),
             str(len(self._highlighted_tasks)),
@@ -1074,8 +1119,11 @@ class TaskProgress(_TaskStateMixin):
 
     def _render_display(self):
         """Render the current display state with caching optimization."""
+        # Get current description (may be dynamic)
+        description = self._get_description()
+
         # Check if we need to re-render
-        current_hash = self._compute_state_hash()
+        current_hash = self._compute_state_hash(description)
         if current_hash == self._last_state_hash and self._render_cache is not None:
             return self._render_cache
 
@@ -1084,20 +1132,20 @@ class TaskProgress(_TaskStateMixin):
         # Build main task line
         if self.main_failed:
             # Split the description to style only the "Failed:" part in red
-            if self.description.startswith("Failed:"):
+            if description.startswith("Failed:"):
                 failed_part = "Failed:"
-                rest_part = self.description[len("Failed:"):].lstrip()
+                rest_part = description[len("Failed:"):].lstrip()
                 main_line = (Text(f"{FAIL_ICON} ", style="red") +
                         Text(failed_part, style="red") +
                         Text(f" {rest_part}", style="default"))
             else:
                 # Fallback if description doesn't start with "Failed:"
-                main_line = Text(f"{FAIL_ICON} ", style="red") + Text(self.description, style="red")
+                main_line = Text(f"{FAIL_ICON} ", style="red") + Text(description, style="red")
         elif self.main_completed:
-            main_line = Text(f"{SUCCESS_ICON} ", style="green") + Text(self.description, style="green")
+            main_line = Text(f"{SUCCESS_ICON} ", style="green") + Text(description, style="green")
         else:
             spinner_text = SPINNER_FRAMES[self.spinner_index]
-            main_line = Text(f"{spinner_text} ", style="magenta") + Text(self.description, style="magenta")
+            main_line = Text(f"{spinner_text} ", style="magenta") + Text(description, style="magenta")
 
         # Build subtask lines
         subtask_lines = self._render_subtask_lines()
@@ -1231,6 +1279,14 @@ class TaskProgress(_TaskStateMixin):
             return "\n".join(lines)
 
     def __enter__(self):
+        # Handle nesting: pause any parent progress
+        self._parent_progress = get_current_progress()
+        if self._parent_progress is not None:
+            self._parent_progress._pause()
+
+        # Set ourselves as the current progress
+        self._context_token = _set_current_progress(self)
+
         if self.leading_newline:
             print()
 
@@ -1246,16 +1302,40 @@ class TaskProgress(_TaskStateMixin):
 
     def _start_spinner(self):
         """Start the spinner animation thread."""
+        self._spinner_paused = False
+
         def spinner_animation():
             while self.live and not self.main_completed and not self.main_failed:
                 time.sleep(SPINNER_UPDATE_INTERVAL)
-                if self.live:
+                if self.live and not self._spinner_paused:
                     self._advance_spinner()
                     self.live.update(self._render_display())
 
         self._spinner_thread = threading.Thread(target=spinner_animation, daemon=True)
         self._spinner_thread.start()
 
+    def _pause(self):
+        """Pause the live display to allow a child progress to render."""
+        self._spinner_paused = True
+        if self.live:
+            # Clear the live display content before stopping so child has a clean slate
+            from rich.text import Text
+            self.live.update(Text(""))
+            self.live.stop()
+            # live.stop() leaves the cursor after the (empty) rendered content.
+            # Move cursor up one line so the child renders in the same place.
+            if not self.is_ci and sys.stdout.isatty():
+                sys.stdout.write("\033[A\r\033[K")
+                sys.stdout.flush()
+
+    def _resume(self):
+        """Resume the live display after a child progress finishes."""
+        self._spinner_paused = False
+        if self.live and not self.main_completed and not self.main_failed:
+            # Re-render current state and restart the live display
+            self.live.update(self._render_display())
+            self.live.start()
+
     def __exit__(self, exc_type, exc_val, exc_tb):
         # Stop timer manager
         self._timer_manager.stop()
@@ -1263,12 +1343,23 @@ class TaskProgress(_TaskStateMixin):
         if exc_type is not None:
             # Exception occurred - show failure message
             self._handle_failure(exc_val)
-            return False  # Don't suppress the exception
+            result = False  # Don't suppress the exception
         else:
             # Success - show completion
             self._handle_success()
+            result = True
+
+        # Restore the parent progress as current
+        if self._context_token is not None:
+            _current_progress.reset(self._context_token)
+            self._context_token = None
+
+        # Resume the parent progress if there was one
+        if self._parent_progress is not None:
+            self._parent_progress._resume()
+            self._parent_progress = None
 
-        return True
+        return result
 
     def _handle_failure(self, exc_val):
         """Handle failure case in context manager exit."""
@@ -1352,7 +1443,7 @@ def _calculate_category_duration(category_name: str, tasks: Dict[str, TaskInfo])
         return sum(task_info.get_duration() for task_info in tasks.values())
 
 
-def create_progress(description: str = "", success_message: str = "", failure_message: str = "",
+def create_progress(description: str | Callable[[], str] = "", success_message: str = "", failure_message: str = "",
 leading_newline: bool = False, trailing_newline: bool = False, show_duration_summary: bool = True):
     """Factory function to create the appropriate progress component based on environment.
 
@@ -1413,14 +1504,20 @@ class NotebookTaskProgress(_TaskStateMixin):
 
     def __init__(
         self,
-        description: str = "",
+        description: str | Callable[[], str] = "",
         success_message: str = "",
         failure_message: str = "",
         leading_newline: bool = False,
         trailing_newline: bool = False,
         show_duration_summary: bool = True,
     ):
-        self.description = description
+        # Description can be a string or callable
+        if callable(description):
+            self._description_fn: Callable[[], str] | None = description
+            self.description = description()  # Initial value
+        else:
+            self._description_fn = None
+            self.description = description
         self.success_message = success_message
         self.failure_message = failure_message
         self.leading_newline = leading_newline
@@ -1447,6 +1544,16 @@ class NotebookTaskProgress(_TaskStateMixin):
         self._current_subtask = ""
         self.busy = False  # Initialize busy state
 
+        # Nesting support
+        self._context_token: Optional[contextvars.Token] = None
+        self._parent_progress: Optional[TaskProgress] = None
+        self._spinner_paused = False
+
+    def _get_description(self) -> str:
+        """Get the current description, calling the function if one was provided."""
+        if self._description_fn is not None:
+            return self._description_fn()
+        return self.description
 
     def _generate_task_id(self) -> str:
         """Generate a unique task ID."""
@@ -1481,15 +1588,18 @@ class NotebookTaskProgress(_TaskStateMixin):
 
         spinner = "⏳⏳⏳⏳" if starting else next(self.spinner_generator)
 
+        # Get current description (may be dynamic)
+        description = self._get_description()
+
         # If there's an active subtask, show ONLY the subtask
         if hasattr(self, '_current_subtask') and self._current_subtask:
             full_message = f"{spinner} {self._current_subtask}"
         else:
             # Otherwise show the main task with subtask count if any
             if len(self._tasks) > 0:
-                full_message = f"{spinner} {self.description} ({len(self._tasks)} active)"
+                full_message = f"{spinner} {description} ({len(self._tasks)} active)"
             else:
-                full_message = f"{spinner} {self.description}"
+                full_message = f"{spinner} {description}"
 
         if len(full_message) > max_width:
             return full_message[:max_width - 3] + "..."
@@ -1539,11 +1649,12 @@ class NotebookTaskProgress(_TaskStateMixin):
 
     def _build_jupyter_lines(self, starting: bool) -> list[str]:
         """Compose the main status and subtasks for Jupyter display."""
+        description = self._get_description()
         if self.busy or starting:
             spinner = SPINNER_FRAMES[0] if starting else next(self.spinner_generator)
-            main_line = f"{spinner} {self.description}"
+            main_line = f"{spinner} {description}"
         else:
-            main_text = self.success_message or self.description
+            main_text = self.success_message or description
             main_line = f"{SUCCESS_ICON} {main_text}"
 
         visible_tasks = self._collect_visible_tasks()
@@ -1582,10 +1693,26 @@ class NotebookTaskProgress(_TaskStateMixin):
     def spinner_task(self):
         """Spinner animation task."""
         while self.busy and self.delay:
-            self.update()
+            if not self._spinner_paused:
+                self.update()
             time.sleep(self.delay) #type: ignore[union-attr] | we only call spinner_task if delay is not None anyway
             self.reset_cursor()
 
+    def _pause(self):
+        """Pause the display to allow a child progress to render."""
+        self._spinner_paused = True
+        # Clear the current line so child can render cleanly
+        if not self.is_jupyter:
+            sys.stdout.write("\r" + " " * DEFAULT_TERMINAL_WIDTH + "\r")
+            sys.stdout.flush()
+
+    def _resume(self):
+        """Resume the display after a child progress finishes."""
+        self._spinner_paused = False
+        # Force an immediate update to restore display
+        if self.busy:
+            self.update()
+
     def _update_subtask_display(self, subtask_text: str):
         """Update sub-task display - shows ONLY the subtask text."""
         # Store the current display state
@@ -1652,6 +1779,14 @@ class NotebookTaskProgress(_TaskStateMixin):
         self._current_display = ""
 
     def __enter__(self):
+        # Handle nesting: pause any parent progress
+        self._parent_progress = get_current_progress()
+        if self._parent_progress is not None:
+            self._parent_progress._pause()
+
+        # Set ourselves as the current progress
+        self._context_token = _set_current_progress(self)
+
         # Skip leading newline for Jupyter - it interferes with IPython display
         if self.leading_newline and not self.is_jupyter:
             rich.print()
@@ -1675,54 +1810,68 @@ class NotebookTaskProgress(_TaskStateMixin):
                 # For Jupyter, IPython display handles formatting
                 if not self.is_jupyter:
                     rich.print(file=sys.stderr)
-                return True
-            return False
-        if self.delay: # will be None for non-interactive environments
-            time.sleep(self.delay)
-
-        # Generate summary BEFORE clearing the spinner line (so we have timing data)
-        # Only generate if show_duration_summary flag is True
-        summary = self.generate_summary() if self.show_duration_summary else ""
-
-        # Clear the spinner line completely
-        self._clear_spinner_line()
-
-        final_message: str | None = None
-        if self.success_message:
-            final_message = f"{SUCCESS_ICON} {self.success_message}"
-        elif summary:
-            final_message = f"{SUCCESS_ICON} Done"
-
-        if final_message:
-            if self.is_jupyter:
-                if self.display is not None:
-                    self.display.update({"text/plain": final_message}, raw=True)
+                result = True
+            else:
+                result = False
+        else:
+            if self.delay: # will be None for non-interactive environments
+                time.sleep(self.delay)
+
+            # Generate summary BEFORE clearing the spinner line (so we have timing data)
+            # Only generate if show_duration_summary flag is True
+            summary = self.generate_summary() if self.show_duration_summary else ""
+
+            # Clear the spinner line completely
+            self._clear_spinner_line()
+
+            final_message: str | None = None
+            if self.success_message:
+                final_message = f"{SUCCESS_ICON} {self.success_message}"
+            elif summary:
+                final_message = f"{SUCCESS_ICON} Done"
+
+            if final_message:
+                if self.is_jupyter:
+                    if self.display is not None:
+                        self.display.update({"text/plain": final_message}, raw=True)
+                    else:
+                        _, display_fn = _load_ipython_display()
+                        self.display = display_fn({"text/plain": final_message}, display_id=True, raw=True)
                 else:
+                    print(final_message)
+            elif self.success_message == "":
+                # When there's no success message, clear the display for notebooks
+                # The summary will be printed below if available
+                if self.is_jupyter:
+                    self.update("")
+                # For non-Jupyter notebooks, _clear_spinner_line() already handled it
+
+            # Print summary if there are completed tasks
+            if summary:
+                if self.is_jupyter:
+                    # Use IPython display to avoid blank stdout lines in notebooks
                     _, display_fn = _load_ipython_display()
-                    self.display = display_fn({"text/plain": final_message}, display_id=True, raw=True)
-            else:
-                print(final_message)
-        elif self.success_message == "":
-            # When there's no success message, clear the display for notebooks
-            # The summary will be printed below if available
-            if self.is_jupyter:
-                self.update("")
-            # For non-Jupyter notebooks, _clear_spinner_line() already handled it
+                    display_fn({"text/plain": summary.strip()}, raw=True)
+                else:
+                    print()
+                    print(summary.strip())  # Summary includes visual separator line
 
-        # Print summary if there are completed tasks
-        if summary:
-            if self.is_jupyter:
-                # Use IPython display to avoid blank stdout lines in notebooks
-                _, display_fn = _load_ipython_display()
-                display_fn({"text/plain": summary.strip()}, raw=True)
-            else:
-                print()
-                print(summary.strip())  # Summary includes visual separator line
+            # Skip trailing newline for Jupyter - it interferes with IPython display
+            if self.trailing_newline and not self.is_jupyter:
+                rich.print()
+            result = True
 
-        # Skip trailing newline for Jupyter - it interferes with IPython display
-        if self.trailing_newline and not self.is_jupyter:
-            rich.print()
-        return True
+        # Restore the parent progress as current
+        if self._context_token is not None:
+            _current_progress.reset(self._context_token)
+            self._context_token = None
+
+        # Resume the parent progress if there was one
+        if self._parent_progress is not None:
+            self._parent_progress._resume()
+            self._parent_progress = None
+
+        return result
 
     def _clear_spinner_line(self):
         """Clear the current spinner line completely."""