diff --git a/.changelog/020.yaml b/.changelog/020.yaml new file mode 100644 index 0000000..60c8b29 --- /dev/null +++ b/.changelog/020.yaml @@ -0,0 +1,8 @@ +name: shell +ts: 2026-06-01 15:20:08.795987+00:00 +type: fix +author: Espen Albert +changelog_message: 'fix: Propagates live_print_scope to shell message_callbacks' +group: shell +message: 'fix: Propagates live_print_scope to shell message_callbacks' +short_sha: 6c245d diff --git a/ask_shell/_internal/_run.py b/ask_shell/_internal/_run.py index 2fae8f1..1720cfa 100644 --- a/ask_shell/_internal/_run.py +++ b/ask_shell/_internal/_run.py @@ -415,10 +415,12 @@ def _execute_run(shell_run: ShellRun) -> ShellRun: config = shell_run.config queue = shell_run._queue + ctx = shell_run._callback_context + def queue_consumer(): for message in queue: try: - shell_run._on_event(message) + ctx.run(shell_run._on_event, message) except BaseException as e: logger.warning(f"Error processing message '{type(message).__name__}' for {shell_run}: {e!r}") logger.exception(e) diff --git a/ask_shell/_internal/_run_test.py b/ask_shell/_internal/_run_test.py index 8faa0e7..e217f57 100644 --- a/ask_shell/_internal/_run_test.py +++ b/ask_shell/_internal/_run_test.py @@ -1,6 +1,8 @@ from rich.ansi import AnsiDecoder -from ask_shell._internal._run import _log_line +from ask_shell._internal._run import _log_line, run_and_wait +from ask_shell._internal.events import ShellRunStdOutput +from ask_shell._internal.live_print_context import get_live_print_context, live_print_scope def test_log_line_strips_ansi_when_enabled(): @@ -8,3 +10,34 @@ def test_log_line_strips_ansi_when_enabled(): raw = "\x1b[32mgreen\x1b[0m\n" assert _log_line(raw, ansi_content=True, decoder=decoder) == "green\n" assert _log_line(raw, ansi_content=False, decoder=decoder) == raw + + +def test_message_callback_inherits_live_print_scope() -> None: + seen: list[str | None] = [] + + def cb(message): + match message: + case ShellRunStdOutput(is_stdout=True): + ctx = get_live_print_context() + seen.append(ctx.prefix if ctx else None) + return False + + with live_print_scope(prefix="[test] "): + run_and_wait("echo ok", message_callbacks=[cb], mute_shell_summary=True) + assert "[test] " in seen + + +def test_message_callback_nested_scope() -> None: + seen: list[str | None] = [] + + def cb(message): + match message: + case ShellRunStdOutput(is_stdout=True): + ctx = get_live_print_context() + seen.append(ctx.prefix if ctx else None) + return False + + with live_print_scope(prefix="outer "): + with live_print_scope(prefix="inner "): + run_and_wait("echo ok", message_callbacks=[cb], mute_shell_summary=True) + assert seen == ["inner "] diff --git a/ask_shell/_internal/models.py b/ask_shell/_internal/models.py index e3349dc..823bec7 100644 --- a/ask_shell/_internal/models.py +++ b/ask_shell/_internal/models.py @@ -1,5 +1,6 @@ from __future__ import annotations +import contextvars import logging import os import platform @@ -293,6 +294,7 @@ class ShellRun: _stderr_log_path: Path | None = field(init=False, default=None) _queue: ShellRunQueueT = field(init=False, default_factory=ClosableQueue) _lock: RLock = field(init=False, default_factory=RLock) + _callback_context: contextvars.Context = field(default_factory=contextvars.copy_context, init=False) @property def has_started(self) -> bool: diff --git a/docs/examples/console/live_print_scope.md b/docs/examples/console/live_print_scope.md index 84d7467..b9b60fa 100644 --- a/docs/examples/console/live_print_scope.md +++ b/docs/examples/console/live_print_scope.md @@ -5,7 +5,7 @@ description: Prefix or suppress live-console scroll lines per scope, including a `live_print_scope` stores per-scope metadata in a [`contextvars`](https://docs.python.org/3/library/contextvars.html) `ContextVar`. [`print_to_live`](../../console/index.md#print_to_live_def) and [`log_to_live`](../../console/index.md#log_to_live_def) read the active scope: they prepend `prefix` to each line and skip output when `suppress` is true. -Use it when parallel work (for example tfdo multi-directory orchestration) needs tagged scroll lines without editing every `print_to_live` call site. Pair it with [`run_pool`](../../shell/run_pool.md): `submit` copies the submitter's context into the worker thread. +Use it when parallel work (for example tfdo multi-directory orchestration) needs tagged scroll lines without editing every `print_to_live` call site. Pair it with [`run_pool`](../../shell/run_pool.md): `submit` copies the submitter's context into the worker thread. With [`run_and_wait`](../../shell/index.md), `message_callbacks` see the scope active when the run was created. ## Read the active scope @@ -74,3 +74,32 @@ with run_pool("demo", total=2, pool_thread_count=2, max_concurrent_submits=2) as print(sorted(results)) #> ['a', 'b'] ``` + +## Context with run_and_wait and message_callbacks + +`ShellRun` snapshots the active context when `run` / `run_and_wait` creates it. The queue consumer runs `message_callbacks` inside that snapshot, so handlers see the same `live_print_scope` as the caller even though dispatch runs on a pool thread. + +```python +from ask_shell._internal._run import run_and_wait +from ask_shell._internal.events import ShellRunStdOutput +from ask_shell.console import get_live_print_context, live_print_scope + +seen: list[str | None] = [] + + +def on_stdout(message): + match message: + case ShellRunStdOutput(is_stdout=True): + ctx = get_live_print_context() + seen.append(ctx.prefix if ctx else None) + return False + + +with live_print_scope(prefix="[dir] "): + run_and_wait("echo ok", message_callbacks=[on_stdout]) + +print(seen) +#> ['[dir] '] +``` + +The snapshot is fixed for the life of that `ShellRun`; changing `live_print_scope` after `run_and_wait` returns does not affect callbacks for an in-flight run. Retries reuse the same snapshot. diff --git a/docs/shell/index.md b/docs/shell/index.md index 189f85c..5b7e26d 100644 --- a/docs/shell/index.md +++ b/docs/shell/index.md @@ -27,7 +27,7 @@ ### class: `ShellRun` -- [source](../../ask_shell/_internal/models.py#L275) +- [source](../../ask_shell/_internal/models.py#L276) > **Since:** 0.3.0 ```python @@ -120,7 +120,7 @@ def kill_all_runs(immediate: bool = False, reason: str = '', abort_timeout: floa ### function: `run_error` -- [source](../../ask_shell/_internal/_run.py#L584) +- [source](../../ask_shell/_internal/_run.py#L586) > **Since:** 0.3.0 ```python @@ -156,7 +156,7 @@ def stop_runs_and_pool(reason: str = 'atexit', immediate: bool = False): ### function: `wait_on_ok_errors` -- [source](../../ask_shell/_internal/_run.py#L591) +- [source](../../ask_shell/_internal/_run.py#L593) > **Since:** 0.3.0 ```python @@ -174,7 +174,7 @@ def wait_on_ok_errors(*runs, timeout: float | None = None, skip_kill_timeouts: b ### exception: `AbortRetryError` -- [source](../../ask_shell/_internal/models.py#L526) +- [source](../../ask_shell/_internal/models.py#L528) - [Example: Stop retrying with a custom error by raising AbortRetryError from should_retry](../examples/shell/AbortRetryError.md) > **Since:** 0.5.0 diff --git a/docs/shell/run.md b/docs/shell/run.md index b6881b4..9a532d4 100644 --- a/docs/shell/run.md +++ b/docs/shell/run.md @@ -2,7 +2,7 @@ ## function: run -- [source](../../ask_shell/_internal/_run.py#L451) +- [source](../../ask_shell/_internal/_run.py#L453) > **Since:** 0.3.0 ```python diff --git a/docs/shell/run_and_wait.md b/docs/shell/run_and_wait.md index e522608..29ee858 100644 --- a/docs/shell/run_and_wait.md +++ b/docs/shell/run_and_wait.md @@ -2,7 +2,7 @@ ## function: run_and_wait -- [source](../../ask_shell/_internal/_run.py#L517) +- [source](../../ask_shell/_internal/_run.py#L519) > **Since:** 0.3.0 ```python diff --git a/docs/shell/shellconfig.md b/docs/shell/shellconfig.md index a500adf..573278a 100644 --- a/docs/shell/shellconfig.md +++ b/docs/shell/shellconfig.md @@ -2,7 +2,7 @@ ## class: ShellConfig -- [source](../../ask_shell/_internal/models.py#L127) +- [source](../../ask_shell/_internal/models.py#L128) > **Since:** 0.3.0 ```python diff --git a/docs/shell/shellerror.md b/docs/shell/shellerror.md index 2572813..0d9d1d0 100644 --- a/docs/shell/shellerror.md +++ b/docs/shell/shellerror.md @@ -2,7 +2,7 @@ ## exception: ShellError -- [source](../../ask_shell/_internal/models.py#L490) +- [source](../../ask_shell/_internal/models.py#L492) > **Since:** 0.3.0 ```python