kash-shell 0.3.20__py3-none-any.whl → 0.3.22__py3-none-any.whl

kash/utils/api_utils/progress_protocol.py

@@ -0,0 +1,299 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Protocol, TypeAlias, TypeVar
+
+T = TypeVar("T")
+TaskID = TypeVar("TaskID")
+
+# Generic task spec types for labeler functions
+TaskSpec = TypeVar("TaskSpec")
+Labeler: TypeAlias = Callable[[int, TaskSpec], str]
+
+# Progress display symbols (consistent with text_styles.py)
+EMOJI_SUCCESS = "[✔︎]"
+EMOJI_FAILURE = "[✘]"
+EMOJI_SKIP = "[-]"
+EMOJI_WARN = "[∆]"
+EMOJI_RETRY = "▵"
+
+
+class TaskState(Enum):
+    """Task execution states."""
+
+    QUEUED = "queued"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+
+
+@dataclass
+class TaskInfo:
+    """Track additional task information beyond basic progress."""
+
+    state: TaskState = TaskState.QUEUED
+    retry_count: int = 0
+    failures: list[str] = field(default_factory=list)
+    label: str = ""
+    total: int = 1
+
+
+@dataclass(frozen=True)
+class TaskSummary:
+    """Summary of task completion states."""
+
+    task_states: list[TaskState]
+
+    @property
+    def queued(self) -> int:
+        """Number of queued tasks."""
+        return sum(1 for state in self.task_states if state == TaskState.QUEUED)
+
+    @property
+    def running(self) -> int:
+        """Number of running tasks."""
+        return sum(1 for state in self.task_states if state == TaskState.RUNNING)
+
+    @property
+    def completed(self) -> int:
+        """Number of completed tasks."""
+        return sum(1 for state in self.task_states if state == TaskState.COMPLETED)
+
+    @property
+    def failed(self) -> int:
+        """Number of failed tasks."""
+        return sum(1 for state in self.task_states if state == TaskState.FAILED)
+
+    @property
+    def skipped(self) -> int:
+        """Number of skipped tasks."""
+        return sum(1 for state in self.task_states if state == TaskState.SKIPPED)
+
+    @property
+    def total(self) -> int:
+        """Total number of tasks."""
+        return len(self.task_states)
+
+    def summary_str(self) -> str:
+        """
+        Generate summary message based on task completion states.
+        """
+        if not self.task_states:
+            return "No tasks to process"
+
+        if self.completed == self.total:
+            return f"All tasks successful: {self.completed}/{self.total} completed"
+        elif self.completed + self.skipped == self.total:
+            return f"All tasks successful: {self.completed}/{self.total} completed, {self.skipped} skipped"
+        elif self.failed == self.total:
+            return f"All tasks failed: {self.failed}/{self.total} failed"
+        else:
+            parts = []
+            if self.completed > 0:
+                parts.append(f"{self.completed}/{self.total} tasks completed")
+            if self.failed > 0:
+                parts.append(f"{self.failed} tasks failed")
+            if self.skipped > 0:
+                parts.append(f"{self.skipped} tasks skipped")
+            if self.queued > 0:
+                parts.append(f"{self.queued} tasks not yet run")
+
+            if self.queued > 0:
+                return "Tasks were interrupted: " + ", ".join(parts)
+            else:
+                return "Tasks had errors: " + ", ".join(parts)
+
+
+class ProgressTracker(Protocol[TaskID]):
+    """
+    Protocol for progress tracking that gather_limited can depend on.
+
+    This allows different implementations (Rich, simple logging, etc.)
+    without creating a hard dependency. Uses a simplified update model.
+    """
+
+    @property
+    def suppress_logs(self) -> bool:
+        """
+        Whether this tracker handles its own display and should suppress
+        standard logging to avoid visual interference.
+        """
+        ...
+
+    async def add(self, label: str, total: int = 1) -> TaskID:
+        """Add a new task to track."""
+        ...
+
+    async def start(self, task_id: TaskID) -> None:
+        """Mark task as started (after rate limiting/queuing)."""
+        ...
+
+    async def update(
+        self,
+        task_id: TaskID,
+        *,
+        progress: int | None = None,
+        label: str | None = None,
+        error_msg: str | None = None,
+    ) -> None:
+        """
+        Update task progress, label, or record a retry attempt.
+
+        Args:
+            task_id: Task ID from add()
+            progress: Steps to advance (None = no change)
+            label: New label (None = no change)
+            error_msg: Error message to record as retry (None = no retry)
+        """
+        ...
+
+    async def finish(
+        self,
+        task_id: TaskID,
+        state: TaskState,
+        message: str = "",
+    ) -> None:
+        """
+        Mark task as finished with final state.
+
+        Args:
+            task_id: Task ID from add()
+            state: Final state (COMPLETED, FAILED, SKIPPED)
+            message: Optional completion/error/skip message
+        """
+        ...
+
+
+class AsyncProgressContext(Protocol[TaskID]):
+    """Protocol for async context manager progress trackers."""
+
+    async def __aenter__(self) -> ProgressTracker[TaskID]:
+        """Start progress tracking."""
+        ...
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
+    ) -> None:
+        """Stop progress tracking."""
+        ...
+
+
+class SimpleProgressTracker:
+    """
+    Basic progress tracker that logs to console.
+    Useful for testing or when Rich is not available.
+    """
+
+    def __init__(self, *, verbose: bool = True, print_fn: Callable[[str], None] = print):
+        self.verbose = verbose
+        self.print_fn = print_fn
+        self._next_id = 1
+        self._tasks: dict[int, TaskInfo] = {}
+
+    @property
+    def suppress_logs(self) -> bool:
+        """Console-based tracker works with standard logging."""
+        return False
+
+    async def add(self, label: str, total: int = 1) -> int:  # pyright: ignore[reportUnusedParameter]
+        task_id = self._next_id
+        self._next_id += 1
+
+        self._tasks[task_id] = TaskInfo(label=label)
+
+        if self.verbose:
+            self.print_fn(f"Queued: {label}")
+
+        return task_id
+
+    async def start(self, task_id: int) -> None:
+        """Mark task as started (after rate limiting/queuing)."""
+        task_info = self._tasks.get(task_id)
+        if not task_info:
+            return
+
+        task_info.state = TaskState.RUNNING
+
+        if self.verbose:
+            self.print_fn(f"Started: {task_info.label}")
+
+    async def update(
+        self,
+        task_id: int,
+        *,
+        progress: int | None = None,  # pyright: ignore[reportUnusedParameter]
+        label: str | None = None,
+        error_msg: str | None = None,
+    ) -> None:
+        task_info = self._tasks.get(task_id)
+        if not task_info:
+            return
+
+        # Update label if provided
+        if label is not None:
+            task_info.label = label
+
+        # Record retry if error message provided
+        if error_msg is not None:
+            task_info.retry_count += 1
+            task_info.failures.append(error_msg)
+
+            if self.verbose:
+                retry_indicator = EMOJI_RETRY * task_info.retry_count
+                self.print_fn(f"Retry {retry_indicator} {task_info.label}: {error_msg}")
+
+    async def finish(
+        self,
+        task_id: int,
+        state: TaskState,
+        message: str = "",
+    ) -> None:
+        task_info = self._tasks.get(task_id)
+        if not task_info:
+            return
+
+        task_info.state = state
+
+        if self.verbose:
+            if state == TaskState.COMPLETED:
+                symbol = EMOJI_SUCCESS
+            elif state == TaskState.FAILED:
+                symbol = EMOJI_FAILURE
+            elif state == TaskState.SKIPPED:
+                symbol = EMOJI_SKIP
+            else:
+                symbol = "?"
+
+            retry_info = (
+                f" (after {task_info.retry_count} retries)" if task_info.retry_count else ""
+            )
+
+            message_info = f": {message}" if message else ""
+            self.print_fn(f"{symbol} {task_info.label}{retry_info}{message_info}")
+
+
+class SimpleProgressContext:
+    """
+    Simple async context manager for SimpleProgressTracker.
+    """
+
+    def __init__(self, *, verbose: bool = True, print_fn: Callable[[str], None] = print):
+        self.verbose = verbose
+        self.print_fn = print_fn
+        self._tracker: SimpleProgressTracker | None = None
+
+    async def __aenter__(self) -> SimpleProgressTracker:
+        self._tracker = SimpleProgressTracker(verbose=self.verbose, print_fn=self.print_fn)
+        return self._tracker
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
+    ) -> None:
+        if self.verbose and self._tracker:
+            # Generate automatic summary
+            task_states = [info.state for info in self._tracker._tasks.values()]
+            summary = TaskSummary(task_states=task_states)
+            self.print_fn(summary.summary_str())

kash/utils/common/function_inspect.py

@@ -11,6 +11,7 @@ from typing import (
     cast,
     get_args,
     get_origin,
+    get_type_hints,
 )
 
 NO_DEFAULT = Parameter.empty  # Alias for clarity
@@ -150,10 +151,19 @@ def inspect_function_params(func: Callable[..., Any], unwrap: bool = True) -> li
     unwrapped_func = inspect.unwrap(func) if unwrap else func
     signature = inspect.signature(unwrapped_func)
 
+    # Get resolved type hints to handle string annotations from __future__ annotations
+    try:
+        type_hints = get_type_hints(unwrapped_func)
+    except (NameError, AttributeError, TypeError):
+        # If we can't resolve type hints (missing imports, etc.), fall back to raw annotations
+        type_hints = {}
+
     param_infos: list[FuncParam] = []
 
     for i, (param_name, param_obj) in enumerate(signature.parameters.items()):
-        effective_type, inner_type, is_optional = _resolve_type_details(param_obj.annotation)
+        # Use resolved type hint if available, otherwise fall back to raw annotation
+        resolved_annotation = type_hints.get(param_name, param_obj.annotation)
+        effective_type, inner_type, is_optional = _resolve_type_details(resolved_annotation)
 
         # Fallback: if type is not resolved from annotation, try to infer from default value.
         if (
@@ -522,3 +532,58 @@ def test_literal_types():
     assert param.name == "flag"
     assert param.effective_type is bool
     assert param.default is True
+
+
+def test_string_annotations():
+    """Test string annotations (from __future__ import annotations) are properly resolved."""
+
+    class Item:  # pyright: ignore[reportUnusedClass]
+        pass
+
+    # Create a function with string annotations to simulate the __future__ annotations behavior
+    def func_with_string_annotations(item):
+        return item
+
+    # Manually set the annotation to a string to simulate __future__ annotations
+    func_with_string_annotations.__annotations__ = {"item": "Item"}
+
+    # This should NOT raise an error and should resolve the type properly
+    params = inspect_function_params(func_with_string_annotations)
+    assert len(params) == 1
+    param = params[0]
+    assert param.name == "item"
+    # The annotation should be preserved as a string
+    assert param.annotation == "Item"
+    # effective_type might be None if "Item" can't be resolved due to scope issues
+    # but the important thing is that it doesn't crash
+    # Note: get_type_hints() can't resolve local classes defined in test functions
+
+    # Test with a known built-in type as a string
+    def func_with_builtin_string_annotation(value):
+        return str(value)
+
+    # Manually set string annotations
+    func_with_builtin_string_annotation.__annotations__ = {"value": "int"}
+
+    params = inspect_function_params(func_with_builtin_string_annotation)
+    assert len(params) == 1
+    param = params[0]
+    assert param.name == "value"
+    assert param.annotation == "int"
+    # This should resolve to the actual int type
+    assert param.effective_type is int
+
+    # Test with a complex type annotation as string
+    def func_with_complex_string_annotation(items):
+        return items
+
+    func_with_complex_string_annotation.__annotations__ = {"items": "list[str]"}
+
+    params = inspect_function_params(func_with_complex_string_annotation)
+    assert len(params) == 1
+    param = params[0]
+    assert param.name == "items"
+    assert param.annotation == "list[str]"
+    # This should resolve to list with inner type str
+    assert param.effective_type is list
+    assert param.inner_type is str

kash/utils/common/testing.py

@@ -3,16 +3,19 @@ from __future__ import annotations
 import os
 from collections.abc import Callable
 from functools import wraps
-from typing import Literal, TypeAlias
+from typing import Literal, ParamSpec, TypeAlias, TypeVar, cast
 
-TestMarker: TypeAlias = Literal["online", "integration"]
+P = ParamSpec("P")
+T = TypeVar("T")
+
+TestMarker: TypeAlias = Literal["online", "integration", "slow"]
 """
 Valid markers for tests. Currently just marking online tests (e.g. LLM APIs that
 that require keys) and more complex integration tests.
 """
 
 
-def enable_if(marker: TestMarker) -> Callable:
+def enable_if(marker: TestMarker) -> Callable[[Callable[P, T]], Callable[P, T]]:
     """
     Mark a test as having external dependencies.
 
@@ -34,13 +37,13 @@ def enable_if(marker: TestMarker) -> Callable:
     ```
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             env_var = f"ENABLE_TESTS_{marker.upper()}"
             if not os.getenv(env_var):
                 print(f"Skipping test function: {func.__name__} (set {env_var}=1 to enable)")
-                return
+                return cast(T, None)
             return func(*args, **kwargs)
 
         # Set pytest markers automatically if pytest is available
@@ -53,6 +56,6 @@ def enable_if(marker: TestMarker) -> Callable:
             # Pytest not available, which is fine for non-test runs
             pass
 
-        return wrapper
+        return wrapper  # type: ignore[return-value]
 
     return decorator