kash-shell 0.3.33__py3-none-any.whl → 0.3.35__py3-none-any.whl

kash/utils/common/import_utils.py

@@ -74,26 +74,108 @@ def import_recursive(
     return tallies
 
 
+def _import_modules_from_package(
+    package: types.ModuleType,
+    package_name: str,
+    max_depth: int = 1,
+    include_private: bool = True,
+    current_depth: int = 0,
+    imported_modules: dict[str, types.ModuleType] | None = None,
+) -> dict[str, types.ModuleType]:
+    """
+    Internal helper to recursively import modules from a package.
+
+    Args:
+        package: The package module to import from
+        package_name: The fully qualified name of the package
+        max_depth: Maximum recursion depth (1 = direct children only)
+        include_private: Whether to import private modules (starting with _)
+        current_depth: Current recursion depth (internal use)
+        imported_modules: Dictionary to accumulate imported modules
+
+    Returns:
+        Dictionary mapping module names to their imported module objects
+    """
+    if imported_modules is None:
+        imported_modules = {}
+
+    if current_depth >= max_depth:
+        return imported_modules
+
+    # Get the module's __path__ if it's a package
+    if not hasattr(package, "__path__"):
+        return imported_modules
+
+    try:
+        for _finder, module_name, ispkg in pkgutil.iter_modules(
+            package.__path__, f"{package_name}."
+        ):
+            # Skip private modules unless requested
+            if not include_private and module_name.split(".")[-1].startswith("_"):
+                continue
+
+            # Skip test modules - they often have special import requirements
+            # and aren't needed for warming the import cache
+            module_parts = module_name.split(".")
+            if any(
+                part in ("tests", "test", "testing", "_test", "_tests") for part in module_parts
+            ):
+                continue
+
+            # Skip already imported modules
+            if module_name in imported_modules:
+                continue
+
+            try:
+                module = importlib.import_module(module_name)
+                imported_modules[module_name] = module
+
+                # Recursively import submodules if it's a package
+                if ispkg and current_depth + 1 < max_depth:
+                    _import_modules_from_package(
+                        module,
+                        module_name,
+                        max_depth=max_depth,
+                        include_private=include_private,
+                        current_depth=current_depth + 1,
+                        imported_modules=imported_modules,
+                    )
+
+            except Exception as e:
+                # Handle various import failures gracefully
+                # This includes ImportError, pytest.Skipped, and other exceptions
+                error_type = type(e).__name__
+                if error_type not in ("ImportError", "AttributeError", "TypeError"):
+                    log.debug(f"  Skipped {module_name}: {error_type}: {e}")
+                # Don't log common/expected import errors to reduce noise
+
+    except Exception as e:
+        log.warning(f"Error iterating modules in {package_name}: {e}")
+
+    return imported_modules
+
+
 def import_namespace_modules(namespace: str) -> dict[str, types.ModuleType]:
     """
     Find and import all modules or packages within a namespace package.
     Returns a dictionary mapping module names to their imported module objects.
     """
-    importlib.import_module(namespace)  # Propagate import errors
+    # Import the main module first
+    main_module = importlib.import_module(namespace)  # Propagate import errors
 
     # Get the package to access its __path__
-    package = sys.modules.get(namespace)
-    if not package or not hasattr(package, "__path__"):
+    if not hasattr(main_module, "__path__"):
         raise ImportError(f"`{namespace}` is not a package or namespace package")
 
-    log.info(f"Discovering modules in `{namespace}` namespace, searching: {package.__path__}")
+    log.info(f"Discovering modules in `{namespace}` namespace, searching: {main_module.__path__}")
+
+    # Use the common helper with depth=1 (no recursion) and include_private=True
+    modules = _import_modules_from_package(
+        main_module, namespace, max_depth=1, include_private=True
+    )
 
-    # Iterate through all modules in the namespace package
-    modules = {}
-    for _finder, module_name, _ispkg in pkgutil.iter_modules(package.__path__, f"{namespace}."):
-        module = importlib.import_module(module_name)  # Propagate import errors
-        log.info(f"Imported module: {module_name} from {module.__file__}")
-        modules[module_name] = module
+    # Add the main module itself
+    modules[namespace] = main_module
 
     log.info(f"Imported {len(modules)} modules from namespace `{namespace}`")
     return modules
@@ -106,8 +188,13 @@ def recursive_reload(
     Recursively reload all modules in the given package that match the filter function.
     Returns a list of module names that were reloaded.
 
-    :param filter_func: A function that takes a module name and returns True if the
-        module should be reloaded.
+    Args:
+        package: The package to reload.
+        filter_func: A function that takes a module name and returns True if the
+            module should be reloaded.
+
+    Returns:
+        List of module names that were reloaded.
     """
     package_name = package.__name__
     modules = {
@@ -124,3 +211,40 @@ def recursive_reload(
         importlib.reload(modules[name])
 
     return module_names
+
+
+def warm_import_library(
+    library_name: str, max_depth: int = 3, include_private: bool = False
+) -> dict[str, types.ModuleType]:
+    """
+    Recursively import all submodules of a library to warm the import cache.
+    This is useful for servers where you want to pay the import cost upfront
+    rather than during request handling.
+
+    Args:
+        library_name: Name of the library to import (e.g., 'litellm', 'openai')
+        max_depth: Maximum depth to recurse into submodules
+        include_private: Whether to import private modules (starting with _)
+
+    Returns:
+        Dictionary mapping module names to their imported module objects
+    """
+    try:
+        # Import the main module first
+        main_module = importlib.import_module(library_name)
+
+        # Use the common helper for recursive imports
+        imported_modules = _import_modules_from_package(
+            main_module, library_name, max_depth=max_depth, include_private=include_private
+        )
+
+        # Add the main module itself
+        imported_modules[library_name] = main_module
+
+    except ImportError as e:
+        log.warning(f"Could not import {library_name}: {e}")
+        return {}
+
+    log.info(f"Warmed {len(imported_modules)} modules from {library_name}")
+
+    return imported_modules

kash/utils/common/s3_utils.py

@@ -1,13 +1,19 @@
 from __future__ import annotations
 
+import os
 import shutil
 import subprocess
+from logging import getLogger
 from pathlib import Path
 
+from dotenv import find_dotenv, load_dotenv
 from sidematter_format.sidematter_format import Sidematter
+from strif import abbrev_str
 
 from kash.utils.common.url import Url, is_s3_url, parse_s3_url
 
+log = getLogger(__name__)
+
 
 def check_aws_cli() -> None:
     """
@@ -19,6 +25,54 @@ def check_aws_cli() -> None:
         )
 
 
+def run_aws_command(cmd: list[str]) -> subprocess.CompletedProcess[str]:
+    """
+    Run an AWS CLI command and capture output.
+    Raises a RuntimeError with stdout/stderr on failure.
+    """
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+        env=os.environ,
+    )
+
+    if result.returncode != 0:
+        # Build a detailed error message
+        error_parts = [f"AWS command failed with exit code {result.returncode}"]
+        error_parts.append(f"Command: {' '.join(cmd)}")
+
+        if result.stdout:
+            error_parts.append(f"stdout: {result.stdout}")
+        if result.stderr:
+            error_parts.append(f"stderr: {result.stderr}")
+
+        raise RuntimeError("\n".join(error_parts))
+
+    return result
+
+
+def reload_aws_env_vars() -> None:
+    """
+    Fresh reload of AWS env vars from .env.local.
+    """
+
+    def aws_creds() -> set[tuple[str, str]]:
+        return {(k, abbrev_str(v, 5)) for k, v in os.environ.items() if k.startswith("AWS_")}
+
+    if len(aws_creds()) == 0:
+        dotenv_path = find_dotenv(".env.local", usecwd=True) or find_dotenv(".env", usecwd=True)
+        load_dotenv(dotenv_path, override=True)
+        if len(aws_creds()) > 0:
+            log.info(
+                "Loaded %s, found AWS credentials: %s",
+                dotenv_path,
+                aws_creds(),
+            )
+        else:
+            log.warning("No AWS credentials found in env or .env files")
+
+
 def get_s3_parent_folder(url: Url) -> Url | None:
     """
     Get the parent folder of an S3 URL, or None if not an S3 URL.
@@ -47,6 +101,7 @@ def s3_sync_to_folder(
     - For a single file: the file URL (and sidematter file/dir URLs if included).
     - For a directory: the destination parent prefix URL (non-recursive reporting).
     """
+    reload_aws_env_vars()
 
     src_path = Path(src_path)
     if not src_path.exists():
@@ -71,7 +126,7 @@ def s3_sync_to_folder(
         for p in sync_paths:
             if p.is_file():
                 # Use sync with include/exclude to leverage default short-circuiting
-                subprocess.run(
+                run_aws_command(
                     [
                         "aws",
                         "s3",
@@ -82,27 +137,54 @@ def s3_sync_to_folder(
                         "*",
                         "--include",
                         p.name,
-                    ],
-                    check=True,
+                    ]
                 )
                 targets.append(Url(dest_prefix + p.name))
             elif p.is_dir():
                 dest_dir = dest_prefix + p.name + "/"
-                subprocess.run(["aws", "s3", "sync", str(p), dest_dir], check=True)
+                run_aws_command(["aws", "s3", "sync", str(p), dest_dir])
                 targets.append(Url(dest_dir))
 
         return targets
     else:
         # Directory mode: sync whole directory.
-        subprocess.run(
+        run_aws_command(
             [
                 "aws",
                 "s3",
                 "sync",
                 str(src_path),
                 dest_prefix,
-            ],
-            check=True,
+            ]
         )
         targets.append(Url(dest_prefix))
         return targets
+
+
+def s3_download_file(s3_url: Url, target_path: str | Path) -> None:
+    """
+    Download a file from S3 to a local path using the AWS CLI.
+
+    Args:
+        s3_url: The S3 URL to download from (s3://bucket/path/to/file)
+        target_path: The local path to save the file to
+    """
+    reload_aws_env_vars()
+
+    if not is_s3_url(s3_url):
+        raise ValueError(f"Source must be an s3:// URL: {s3_url}")
+
+    check_aws_cli()
+
+    target_path = Path(target_path)
+
+    # Use aws s3 cp to download the file
+    run_aws_command(
+        [
+            "aws",
+            "s3",
+            "cp",
+            str(s3_url),
+            str(target_path),
+        ]
+    )

kash/utils/rich_custom/multitask_status.py

@@ -72,6 +72,8 @@ RUNNING_SYMBOL = ""
 DEFAULT_LABEL_WIDTH = 40
 DEFAULT_PROGRESS_WIDTH = 20
 
+MAX_DISPLAY_TASKS = 20
+
 
 # Calculate spinner width to maintain column alignment
 def _get_spinner_width(spinner_name: str) -> int:
@@ -101,6 +103,9 @@ class StatusSettings:
     transient: bool = True
     refresh_per_second: float = 10
     styles: StatusStyles = DEFAULT_STYLES
+    # Maximum number of tasks to keep visible in the live display.
+    # Older completed/skipped/failed tasks beyond this cap will be removed from the live view.
+    max_display_tasks: int = MAX_DISPLAY_TASKS
 
 
 class SpinnerStatusColumn(ProgressColumn):
@@ -298,6 +303,10 @@ class MultiTaskStatus(AbstractAsyncContextManager):
         self._task_info: dict[int, TaskInfo] = {}
         self._next_id: int = 1
         self._rich_task_ids: dict[int, TaskID] = {}  # Map our IDs to Rich Progress IDs
+        # Track order of tasks added to the Progress so we can prune oldest completed ones
+        self._displayed_task_order: list[int] = []
+        # Track tasks pruned from the live display so we don't re-add them later
+        self._pruned_task_ids: set[int] = set()
 
         # Unified live integration
         self._unified_live: Any | None = None  # Reference to the global unified live
@@ -442,6 +451,10 @@ class MultiTaskStatus(AbstractAsyncContextManager):
                 progress_display=None,
             )
             self._rich_task_ids[task_id] = rich_task_id
+            self._displayed_task_order.append(task_id)
+
+            # Prune if too many tasks are visible (prefer removing completed ones)
+            self._prune_completed_tasks_if_needed()
 
     async def set_progress_display(self, task_id: int, display: RenderableType) -> None:
         """
@@ -536,18 +549,31 @@ class MultiTaskStatus(AbstractAsyncContextManager):
 
             # Complete the progress bar and stop spinner
             if rich_task_id is not None:
-                total = self._progress.tasks[rich_task_id].total or 1
+                # Safely find the Task by id; Progress.tasks is a list, not a dict
+                task_obj = next((t for t in self._progress.tasks if t.id == rich_task_id), None)
+                if task_obj is not None and task_obj.total is not None:
+                    total = task_obj.total
+                else:
+                    total = task_info.steps_total or 1
                 self._progress.update(rich_task_id, completed=total, task_info=task_info)
             else:
-                # Task was never started, but we still need to add it to show completion
-                rich_task_id = self._progress.add_task(
-                    "",
-                    total=task_info.steps_total,
-                    label=task_info.label,
-                    completed=task_info.steps_total,
-                    task_info=task_info,
-                )
-                self._rich_task_ids[task_id] = rich_task_id
+                # If this task was pruned from the live display, skip re-adding it
+                if task_id in self._pruned_task_ids:
+                    pass
+                else:
+                    # Task was never started; add a completed row so it appears once
+                    rich_task_id = self._progress.add_task(
+                        "",
+                        total=task_info.steps_total,
+                        label=task_info.label,
+                        completed=task_info.steps_total,
+                        task_info=task_info,
+                    )
+                    self._rich_task_ids[task_id] = rich_task_id
+                    self._displayed_task_order.append(task_id)
+
+            # After finishing, prune completed tasks to respect max visible cap
+            self._prune_completed_tasks_if_needed()
 
     def get_task_info(self, task_id: int) -> TaskInfo | None:
         """Get additional task information."""
@@ -567,6 +593,54 @@ class MultiTaskStatus(AbstractAsyncContextManager):
         """Get console instance for additional output above progress."""
         return self._progress.console
 
+    def _prune_completed_tasks_if_needed(self) -> None:
+        """
+        Ensure at most `max_display_tasks` tasks are visible by removing the oldest
+        completed/skipped/failed tasks first. Running or waiting tasks are never
+        removed by this method.
+        Note: This method assumes it's called under self._lock.
+        """
+        max_visible = self.settings.max_display_tasks
+
+        # Nothing to prune or unlimited
+        if max_visible <= 0:
+            return
+
+        # Count visible tasks (those with a Rich task id present)
+        visible_task_ids = [tid for tid in self._displayed_task_order if tid in self._rich_task_ids]
+        excess = len(visible_task_ids) - max_visible
+        if excess <= 0:
+            return
+
+        # Build list of terminal tasks that can be pruned (oldest first)
+        terminal_tasks = []
+        for tid in self._displayed_task_order:
+            if tid not in self._rich_task_ids:
+                continue
+            info = self._task_info.get(tid)
+            if info and info.state in (
+                TaskState.COMPLETED,
+                TaskState.FAILED,
+                TaskState.SKIPPED,
+            ):
+                terminal_tasks.append(tid)
+
+        # Remove the oldest terminal tasks up to the excess count
+        tasks_to_remove = terminal_tasks[:excess]
+
+        for tid in tasks_to_remove:
+            rich_tid = self._rich_task_ids.pop(tid, None)
+            if rich_tid is not None:
+                # Remove from Rich progress display
+                self._progress.remove_task(rich_tid)
+            # Mark as pruned so we don't re-add on finish
+            self._pruned_task_ids.add(tid)
+
+        # Efficiently rebuild the displayed task order without the removed tasks
+        self._displayed_task_order = [
+            tid for tid in self._displayed_task_order if tid not in tasks_to_remove
+        ]
+
 
 ## Tests
 

kash/utils/text_handling/markdown_footnotes.py

@@ -1,48 +1,19 @@
 from __future__ import annotations
 
-import re
 from dataclasses import dataclass, field
 from typing import Any
 
-from flowmark import flowmark_markdown, line_wrap_by_sentence
 from marko import Markdown
+from marko.block import Document
 from marko.ext import footnote
 
-from kash.utils.text_handling.markdown_utils import comprehensive_transform_tree
-
-
-def _normalize_footnotes_in_markdown(content: str) -> str:
-    """
-    Ensure blank lines between consecutive footnote definitions.
-
-    Marko has a bug where consecutive footnotes without blank lines are parsed
-    as a single footnote. This adds blank lines where needed.
-    """
-    lines = content.split("\n")
-    result = []
-    i = 0
-
-    while i < len(lines):
-        line = lines[i]
-        result.append(line)
-
-        # Check if this is a footnote definition
-        if re.match(r"^\[\^[^\]]+\]:", line):
-            # Look ahead to see if the next non-empty line is also a footnote
-            j = i + 1
-            while j < len(lines) and not lines[j].strip():
-                result.append(lines[j])
-                j += 1
-
-            if j < len(lines) and re.match(r"^\[\^[^\]]+\]:", lines[j]):
-                # Next non-empty line is also a footnote, add blank line
-                result.append("")
-
-            i = j
-        else:
-            i += 1
-
-    return "\n".join(result)
+from kash.utils.text_handling.markdown_utils import (
+    MARKDOWN as DEFAULT_MARKDOWN,
+)
+from kash.utils.text_handling.markdown_utils import (
+    comprehensive_transform_tree,
+    normalize_footnotes_in_markdown,
+)
 
 
 @dataclass
@@ -81,15 +52,17 @@ class MarkdownFootnotes:
             MarkdownFootnotes instance with all footnotes indexed by ID
         """
         if markdown_parser is None:
-            markdown_parser = flowmark_markdown(line_wrap_by_sentence(is_markdown=True))
+            markdown_parser = DEFAULT_MARKDOWN
 
         # Normalize to work around marko bug with consecutive footnotes
-        normalized_content = _normalize_footnotes_in_markdown(content)
+        normalized_content = normalize_footnotes_in_markdown(content)
         document = markdown_parser.parse(normalized_content)
         return MarkdownFootnotes.from_document(document, markdown_parser)
 
     @staticmethod
-    def from_document(document: Any, markdown_parser: Markdown | None = None) -> MarkdownFootnotes:
+    def from_document(
+        document: Document, markdown_parser: Markdown | None = None
+    ) -> MarkdownFootnotes:
         """
         Extract all footnotes from a parsed markdown document.
 
@@ -102,7 +75,7 @@ class MarkdownFootnotes:
             MarkdownFootnotes instance with all footnotes indexed by ID
         """
         if markdown_parser is None:
-            markdown_parser = flowmark_markdown(line_wrap_by_sentence(is_markdown=True))
+            markdown_parser = DEFAULT_MARKDOWN
 
         footnotes_dict: dict[str, FootnoteInfo] = {}
 
@@ -206,9 +179,9 @@ def extract_footnote_references(content: str, markdown_parser: Markdown | None =
         List of unique footnote IDs that are referenced (with the ^)
     """
     if markdown_parser is None:
-        markdown_parser = flowmark_markdown(line_wrap_by_sentence(is_markdown=True))
+        markdown_parser = DEFAULT_MARKDOWN
 
-    normalized_content = _normalize_footnotes_in_markdown(content)
+    normalized_content = normalize_footnotes_in_markdown(content)
     document = markdown_parser.parse(normalized_content)
     references: list[str] = []
     seen: set[str] = set()