ask-shell 0.5.2__tar.gz → 0.6.0__tar.gz

{ask_shell-0.5.2 → ask_shell-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ask-shell
-Version: 0.5.2
+Version: 0.6.0
 Summary: CLIs with prompts, shell runs, and testable flows
 Author-email: EspenAlbert <espen.albert1@gmail.com>
 License-Expression: MIT

{ask_shell-0.5.2 → ask_shell-0.6.0}/ask_shell/__init__.py

@@ -6,7 +6,7 @@ from ask_shell import shell_events
 from ask_shell import ask
 from ask_shell.settings import AskShellSettings
 
-VERSION = "0.5.2"
+VERSION = "0.6.0"
 __all__ = [
     "shell",
     "console",

ask_shell-0.6.0/ask_shell/_internal/run_pool.py

@@ -0,0 +1,121 @@
+import logging
+import time
+from concurrent.futures import Future, ThreadPoolExecutor
+from concurrent.futures import wait as futures_wait
+from dataclasses import dataclass, field
+from math import ceil
+from threading import RLock
+from typing import Any, Callable, Protocol, TypeVar
+
+from ask_shell._internal._run import (
+    THREADS_PER_RUN,
+    get_pool,
+    handle_interrupt_wait,
+    max_run_count_for_workers,
+    wait_if_many_runs,
+)
+from ask_shell._internal.rich_progress import new_task
+from ask_shell.settings import AskShellSettings
+
+logger = logging.getLogger(__name__)
+T_co = TypeVar("T_co", covariant=True)
+
+
+class SubmitFunc(Protocol[T_co]):
+    def __call__(self, *args: Any, **kwargs: Any) -> T_co: ...
+
+
+@dataclass
+class run_pool:
+    task_name: str
+    total: int = 0
+    max_concurrent_submits: int = field(default=4)
+    threads_used_per_submit: int = (
+        THREADS_PER_RUN + 1
+    )  # If you are using `run` or `run_and_wait` this should be `THREADS_PER_RUN` + extra threads for your own tasks
+    pool_thread_count: int | None = None
+    sleep_time: float = 1
+    sleep_callback: Callable[[], Any] | None = None
+    exit_wait_timeout: float | None = (
+        None  # If set, will wait for the pool to finish before exiting the context manager
+    )
+
+    pool: ThreadPoolExecutor = field(init=False)
+    _owns_pool: bool = field(init=False, default=False)
+    _pool_max_workers: int = field(init=False)
+    _max_run_count_with_this_pool: int = field(init=False)
+    _lock: RLock = field(init=False, default_factory=RLock)
+    _pending_count: int = field(init=False, default=0)
+    _task: new_task | None = field(init=False, default=None)
+    _futures: list[Future] = field(init=False, default_factory=list)
+
+    def __post_init__(self):
+        if self.pool_thread_count is not None:
+            self.pool = ThreadPoolExecutor(max_workers=self.pool_thread_count)
+            self._owns_pool = True
+            self._pool_max_workers = self.pool_thread_count
+            # dedicated pool: each concurrent submit reserves 1 run slot on the global pool
+            runs_needed = self.max_concurrent_submits
+            runs_available = max_run_count_for_workers()
+            logger.debug(
+                f"run_pool '{self.task_name}': dedicated pool with {self.pool_thread_count} workers, "
+                f"global pool reserves {runs_needed}/{runs_available} run slots"
+            )
+        else:
+            self.pool = get_pool()
+            self._pool_max_workers = self.pool._max_workers
+            # shared pool: submits + their shell runs share the same threads
+            workers_at_full_load = self.max_concurrent_submits * self.threads_used_per_submit
+            runs_needed = ceil(workers_at_full_load / THREADS_PER_RUN)
+            runs_available = max_run_count_for_workers(self._pool_max_workers)
+
+        assert runs_needed < runs_available, (
+            f"Run slots needed ({runs_needed}) exceed capacity ({runs_available}). "
+            f"Adjust {AskShellSettings.ENV_NAME_THREAD_COUNT} or decrease `max_concurrent_submits`."
+        )
+        self._max_run_count_with_this_pool = runs_available - runs_needed
+
+    def _on_submit_done(self, _future: Future):
+        with self._lock:
+            self._pending_count -= 1
+            if task := self._task:
+                task.update(advance=1)
+
+    def submit(self, fn: SubmitFunc[T_co], /, *args, **kwargs) -> Future[T_co]:
+        """Submit a task to the pool. Blocks if max_concurrent_submits are already in flight."""
+        with self._lock:
+            self._pending_count += 1
+        with handle_interrupt_wait(interrupt_message=f"run_pool submit for {self.task_name}"):
+            while self._pending_count > self.max_concurrent_submits:
+                if self.sleep_callback:
+                    self.sleep_callback()
+                time.sleep(self.sleep_time)
+        wait_if_many_runs(
+            max_run_count=self._max_run_count_with_this_pool,
+            sleep_time=self.sleep_time,
+            sleep_callback=self.sleep_callback,
+        )
+        future = self.pool.submit(fn, *args, **kwargs)
+        future.add_done_callback(self._on_submit_done)
+        with self._lock:
+            self._futures.append(future)
+        return future
+
+    def __enter__(self):
+        self._task = new_task(self.task_name, self.total)
+        self._task.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        with self._lock:
+            futures = list(self._futures)
+        if futures:
+            with handle_interrupt_wait(interrupt_message=f"interrupt in `run_pool` exit method for {self.task_name}"):
+                futures_wait(futures, timeout=self.exit_wait_timeout)
+            with self._lock:
+                self._futures.clear()
+
+        if self._owns_pool:
+            self.pool.shutdown(wait=True)
+        if task := self._task:
+            task.__exit__(exc_type, exc_value, traceback)

{ask_shell-0.5.2 → ask_shell-0.6.0}/ask_shell/_internal/typer_command.py

@@ -88,6 +88,42 @@ def track_progress_decorator(
     return decorator
 
 
+def _wrap_typer_tree_commands(
+    app: typer.Typer,
+    *,
+    settings: AskShellSettings,
+    log_path_prefix: str,
+    skip_except_hook: bool,
+    use_app_name_command_for_logs: bool,
+    render_rich_error_on_sys_exit: bool,
+) -> None:
+    for command in app.registered_commands:
+        command.callback = track_progress_decorator(
+            skip_except_hook=skip_except_hook,
+            settings=settings,
+            use_app_name_command_for_logs=use_app_name_command_for_logs,
+            app_name=log_path_prefix,
+            command_name=command.name or command.callback.__name__,  # type: ignore
+            skip_rich_exception=not render_rich_error_on_sys_exit,
+        )(
+            command.callback  # type: ignore
+        )
+    for group in app.registered_groups:
+        nested = group.typer_instance
+        if nested is None:
+            continue
+        segment = group.name or nested.info.name or "group"
+        child_prefix = f"{log_path_prefix}/{segment}"
+        _wrap_typer_tree_commands(
+            nested,
+            settings=settings,
+            log_path_prefix=child_prefix,
+            skip_except_hook=skip_except_hook,
+            use_app_name_command_for_logs=use_app_name_command_for_logs,
+            render_rich_error_on_sys_exit=render_rich_error_on_sys_exit,
+        )
+
+
 def remove_secrets(message: str, secrets: list[str]) -> str:
     for secret in secrets:
         message = message.replace(secret, "***")
@@ -141,18 +177,15 @@ def configure_logging(
     render_rich_error_on_sys_exit: bool = False,
 ) -> logging.Handler:
     settings = settings or AskShellSettings.from_env()
-    app_name = app.info.name or "typer_app"
-    for command in app.registered_commands:
-        command.callback = track_progress_decorator(
-            skip_except_hook=skip_except_hook,
-            settings=settings,
-            use_app_name_command_for_logs=use_app_name_command_for_logs,
-            app_name=app_name,
-            command_name=command.name or command.callback.__name__,  # type: ignore
-            skip_rich_exception=not render_rich_error_on_sys_exit,
-        )(
-            command.callback  # type: ignore
-        )
+    root_prefix = app.info.name or "typer_app"
+    _wrap_typer_tree_commands(
+        app,
+        settings=settings,
+        log_path_prefix=root_prefix,
+        skip_except_hook=skip_except_hook,
+        use_app_name_command_for_logs=use_app_name_command_for_logs,
+        render_rich_error_on_sys_exit=render_rich_error_on_sys_exit,
+    )
     handler = RichHandler(rich_tracebacks=False, level=settings.log_level, console=get_live_console())
     logging.basicConfig(
         level=settings.log_level,

{ask_shell-0.5.2 → ask_shell-0.6.0}/ask_shell/settings.py

@@ -102,7 +102,7 @@ class AskShellSettings(StaticSettings):
     )
     ENV_NAME_THREAD_COUNT: ClassVar[str] = f"{ENV_PREFIX}THREAD_COUNT"
     thread_count: int = Field(
-        default=50,
+        default=100,
         alias=ENV_NAME_THREAD_COUNT,
         description="Thread count for ask-shell pool",
     )

{ask_shell-0.5.2 → ask_shell-0.6.0}/pyproject.toml

@@ -2,7 +2,7 @@
 
 [project]
 name = "ask-shell"
-version = "0.5.2"
+version = "0.6.0"
 description = "CLIs with prompts, shell runs, and testable flows"
 requires-python = ">=3.13"
 license = "MIT"

ask_shell-0.5.2/ask_shell/_internal/run_pool.py

@@ -1,100 +0,0 @@
-import time
-from concurrent.futures import Future, ThreadPoolExecutor
-from dataclasses import dataclass, field
-from math import ceil
-from threading import Event, RLock
-from typing import Any, Callable, Protocol, TypeVar
-
-from ask_shell._internal._run import (
-    THREADS_PER_RUN,
-    get_pool,
-    handle_interrupt_wait,
-    max_run_count_for_workers,
-    wait_if_many_runs,
-)
-from ask_shell._internal.rich_progress import new_task
-from ask_shell.settings import AskShellSettings
-
-T_co = TypeVar("T_co", covariant=True)
-
-
-class SubmitFunc(Protocol[T_co]):
-    def __call__(self, *args: Any, **kwargs: Any) -> T_co: ...
-
-
-@dataclass
-class run_pool:
-    task_name: str
-    total: int = 0
-    max_concurrent_submits: int = field(default=4)
-    threads_used_per_submit: int = (
-        THREADS_PER_RUN + 1
-    )  # If you are using `run` or `run_and_wait` this should be `THREADS_PER_RUN` + extra threads for your own tasks
-    sleep_time: float = 1
-    sleep_callback: Callable[[], Any] | None = None
-    exit_wait_timeout: float | None = (
-        None  # If set, will wait for the pool to finish before exiting the context manager
-    )
-
-    pool: ThreadPoolExecutor = field(init=False, default_factory=get_pool)
-    _pool_max_workers: int = field(init=False)
-    _max_run_count_with_this_pool: int = field(init=False)
-    _lock: RLock = field(init=False, default_factory=RLock)
-    _current_submit_count: int = field(init=False, default=0)
-    _task: new_task | None = field(init=False, default=None)
-    _event: Event = field(init=False, default_factory=Event)
-
-    def __post_init__(self):
-        self._pool_max_workers = self.pool._max_workers
-        max_run_count = max_run_count_for_workers(self._pool_max_workers)
-        workers_required_if_full = self.max_concurrent_submits * self.threads_used_per_submit
-        run_count_used_by_this_pool = ceil(workers_required_if_full / THREADS_PER_RUN)
-        assert run_count_used_by_this_pool < max_run_count, (
-            f"Run count used by this pool ({run_count_used_by_this_pool}) exceeds max run count ({max_run_count}). Adjust {AskShellSettings.ENV_NAME_THREAD_COUNT} environment variable or decrease `max_concurrent_submits` parameter."
-        )
-        self._max_run_count_with_this_pool = max_run_count - run_count_used_by_this_pool
-
-    def _on_submit_done(self):
-        """Callback to be called when a submit is done. This is used to decrement the current submit count."""
-        with self._lock:
-            self._current_submit_count -= 1
-            if task := self._task:
-                task.update(advance=1)
-            if self._current_submit_count == 0:
-                self._event.set()
-
-    def submit(self, fn: SubmitFunc[T_co], /, *args, **kwargs) -> Future[T_co]:
-        """Submit a task to the pool. Might block if the pool is full."""
-
-        # problem: There is a bit of lag from submit until the run is actually started,
-        with self._lock:
-            self._current_submit_count += 1
-            if self._current_submit_count == 1:
-                self._event = Event()  # reset the event when the first submit is made
-        with handle_interrupt_wait(interrupt_message=f"run_pool submit for {self.task_name}"):
-            while self._current_submit_count >= self.max_concurrent_submits:
-                if self.sleep_callback:
-                    self.sleep_callback()
-                time.sleep(self.sleep_time)
-        # in case more runs are already submitted
-        wait_if_many_runs(
-            max_run_count=self._max_run_count_with_this_pool,
-            sleep_time=self.sleep_time,
-            sleep_callback=self.sleep_callback,
-        )
-        future = self.pool.submit(fn, *args, **kwargs)
-        future.add_done_callback(lambda _: self._on_submit_done())
-        return future
-
-    def __enter__(self):
-        self._task = new_task(self.task_name, self.total)
-        self._task.__enter__()
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback) -> None:
-        # no cleanup necessary, the pool will be cleaned up automatically due to atexit call
-        with handle_interrupt_wait(interrupt_message=f"interrupt in `run_pool` exit method for {self.task_name}"):
-            self._event.wait(self.exit_wait_timeout)
-
-        if task := self._task:
-            task.__exit__(exc_type, exc_value, traceback)