diff --git a/.changelog/019.yaml b/.changelog/019.yaml new file mode 100644 index 0000000..c509830 --- /dev/null +++ b/.changelog/019.yaml @@ -0,0 +1,20 @@ +name: get_live_print_context +ts: 2026-06-01 12:26:33.762426+00:00 +type: make_public +details: created in _internal/live_print_context.py +full_path: _internal.live_print_context.get_live_print_context +group: console +--- +name: live_print_scope +ts: 2026-06-01 12:26:33.764529+00:00 +type: make_public +details: created in _internal/live_print_context.py +full_path: _internal.live_print_context.live_print_scope +group: console +--- +name: LivePrintContext +ts: 2026-06-01 12:26:37.231694+00:00 +type: make_public +details: created in _internal/live_print_context.py +full_path: _internal.live_print_context.LivePrintContext +group: console diff --git a/.groups.yaml b/.groups.yaml index 0e4cb6c..bc4ff08 100644 --- a/.groups.yaml +++ b/.groups.yaml @@ -29,12 +29,16 @@ groups: - name: console owned_modules: - _internal._run_env + - _internal.live_print_context - _internal.rich_live - _internal.rich_progress - _internal.typer_command owned_refs: - _internal._run_env.disable_interactive_shell - _internal._run_env.interactive_shell + - _internal.live_print_context.LivePrintContext + - _internal.live_print_context.get_live_print_context + - _internal.live_print_context.live_print_scope - _internal.rich_live.RemoveLivePart - _internal.rich_live.add_renderable - _internal.rich_live.get_live_console diff --git a/ask_shell/_internal/live_print_context.py b/ask_shell/_internal/live_print_context.py new file mode 100644 index 0000000..6620475 --- /dev/null +++ b/ask_shell/_internal/live_print_context.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Iterator + + +@dataclass +class LivePrintContext: + prefix: str = "" + suppress: bool = False + + +_live_print_ctx: ContextVar[LivePrintContext | None] = ContextVar("live_print_ctx", default=None) + + +def get_live_print_context() -> LivePrintContext | None: + return _live_print_ctx.get() + + +@contextmanager +def live_print_scope(*, prefix: str = "", suppress: bool = False) -> Iterator[None]: + token = _live_print_ctx.set(LivePrintContext(prefix=prefix, suppress=suppress)) + try: + yield + finally: + _live_print_ctx.reset(token) diff --git a/ask_shell/_internal/live_print_context_test.py b/ask_shell/_internal/live_print_context_test.py new file mode 100644 index 0000000..22a7ec3 --- /dev/null +++ b/ask_shell/_internal/live_print_context_test.py @@ -0,0 +1,37 @@ +from unittest.mock import MagicMock, patch + +from ask_shell._internal.live_print_context import live_print_scope +from ask_shell._internal.rich_live import log_to_live, print_to_live + + +@patch("ask_shell._internal.rich_live.get_live_console") +def test_live_print_scope_prefix(mock_get_console: MagicMock) -> None: + console = MagicMock() + mock_get_console.return_value = console + with live_print_scope(prefix="[dir] "): + print_to_live("x") + console.print.assert_called_once() + assert console.print.call_args[0][0] == "[dir] " + + +@patch("ask_shell._internal.rich_live.get_live_console") +def test_live_print_scope_suppress(mock_get_console: MagicMock) -> None: + console = MagicMock() + mock_get_console.return_value = console + with live_print_scope(suppress=True): + print_to_live("x") + log_to_live("y") + console.print.assert_not_called() + console.log.assert_not_called() + + +@patch("ask_shell._internal.rich_live.get_live_console") +def test_nested_live_print_scope(mock_get_console: MagicMock) -> None: + console = MagicMock() + mock_get_console.return_value = console + with live_print_scope(prefix="outer "): + print_to_live("a") + with live_print_scope(prefix="inner "): + print_to_live("b") + print_to_live("c") + assert [c[0][0] for c in console.print.call_args_list] == ["outer ", "inner ", "outer "] diff --git a/ask_shell/_internal/rich_live.py b/ask_shell/_internal/rich_live.py index fc571e0..6667e68 100644 --- a/ask_shell/_internal/rich_live.py +++ b/ask_shell/_internal/rich_live.py @@ -14,6 +14,7 @@ from zero_3rdparty.id_creator import simple_id from ask_shell._internal._run_env import resolve_terminal_dimensions +from ask_shell._internal.live_print_context import get_live_print_context _live: Live | None = None _lock = RLock() @@ -180,6 +181,17 @@ def get_live_console() -> Console: return get_live().console +def _objects_for_live_print(*objects: Any) -> tuple[Any, ...] | None: + ctx = get_live_print_context() + if ctx is None: + return objects + if ctx.suppress: + return None + if ctx.prefix: + return (ctx.prefix, *objects) + return objects + + def print_to_live( *objects: Any, sep: str = " ", @@ -197,8 +209,11 @@ def print_to_live( soft_wrap: bool | None = None, new_line_start: bool = False, ): + live_objects = _objects_for_live_print(*objects) + if live_objects is None: + return get_live_console().print( - *objects, + *live_objects, sep=sep, end=end, style=style, @@ -228,8 +243,11 @@ def log_to_live( log_locals: bool = False, _stack_offset: int = 1, ) -> None: + live_objects = _objects_for_live_print(*objects) + if live_objects is None: + return get_live_console().log( - *objects, + *live_objects, sep=sep, end=end, style=style, diff --git a/ask_shell/_internal/run_pool.py b/ask_shell/_internal/run_pool.py index 08f6597..17d9dc3 100644 --- a/ask_shell/_internal/run_pool.py +++ b/ask_shell/_internal/run_pool.py @@ -1,3 +1,4 @@ +import contextvars import logging import time from concurrent.futures import Future, ThreadPoolExecutor @@ -95,7 +96,8 @@ def submit(self, fn: SubmitFunc[T_co], /, *args, **kwargs) -> Future[T_co]: sleep_time=self.sleep_time, sleep_callback=self.sleep_callback, ) - future = self.pool.submit(fn, *args, **kwargs) + ctx = contextvars.copy_context() + future = self.pool.submit(ctx.run, fn, *args, **kwargs) future.add_done_callback(self._on_submit_done) with self._lock: self._futures.append(future) diff --git a/ask_shell/_internal/run_pool_test.py b/ask_shell/_internal/run_pool_test.py index bf97933..e98722a 100644 --- a/ask_shell/_internal/run_pool_test.py +++ b/ask_shell/_internal/run_pool_test.py @@ -1,8 +1,9 @@ from concurrent.futures import Future, ThreadPoolExecutor -from threading import Thread +from threading import Barrier, Thread from unittest.mock import MagicMock, patch from ask_shell._internal._run import get_pool, max_run_count_for_workers, wait_if_many_runs +from ask_shell._internal.live_print_context import get_live_print_context, live_print_scope from ask_shell._internal.rich_progress import new_task from ask_shell._internal.run_pool import run_pool @@ -87,7 +88,7 @@ def task_fn(): mock_pool = MagicMock(spec=ThreadPoolExecutor) mock_pool._max_workers = 20 - mock_pool.submit = MagicMock(side_effect=lambda fn, *a, **kw: _immediate_future(fn, *a, **kw)) + mock_pool.submit = MagicMock(side_effect=_submit_via_context_run) with patch(f"{_module}.{get_pool.__name__}", return_value=mock_pool): rp = run_pool(task_name="single", total=1, max_concurrent_submits=1, threads_used_per_submit=4, sleep_time=0.01) @@ -121,7 +122,7 @@ def task_fn(index: int): mock_pool = MagicMock(spec=ThreadPoolExecutor) mock_pool._max_workers = 20 - mock_pool.submit = MagicMock(side_effect=lambda fn, *a, **kw: _immediate_future(fn, *a, **kw)) + mock_pool.submit = MagicMock(side_effect=_submit_via_context_run) with patch(f"{_module}.{get_pool.__name__}", return_value=mock_pool): rp = run_pool(task_name="multi", total=3, max_concurrent_submits=3, threads_used_per_submit=4, sleep_time=0.01) @@ -146,8 +147,30 @@ def do_submits(): assert sorted(results) == [0, 1, 2] -def _immediate_future(fn, *args, **kwargs): - """Run fn synchronously and return a resolved Future with done callbacks.""" +def test_run_pool_submit_copies_context() -> None: + results: list[str] = [] + barrier = Barrier(2) + + def _read_prefix() -> None: + barrier.wait() + ctx = get_live_print_context() + results.append(ctx.prefix if ctx else "") + + with run_pool("test", total=2, pool_thread_count=2, max_concurrent_submits=2) as pool: + with live_print_scope(prefix="a"): + f_a = pool.submit(_read_prefix) + with live_print_scope(prefix="b"): + f_b = pool.submit(_read_prefix) + f_a.result() + f_b.result() + assert sorted(results) == ["a", "b"] + + +def _submit_via_context_run(run_fn, fn, /, *args, **kwargs): + return _immediate_future(run_fn, fn, *args, **kwargs) + + +def _immediate_future(fn, /, *args, **kwargs): fut = Future() try: result = fn(*args, **kwargs) diff --git a/ask_shell/console.py b/ask_shell/console.py index 3d78fb8..f3419cd 100644 --- a/ask_shell/console.py +++ b/ask_shell/console.py @@ -1,6 +1,9 @@ # Generated by pkg-ext from ask_shell._internal._run_env import disable_interactive_shell as _disable_interactive_shell from ask_shell._internal._run_env import interactive_shell as _interactive_shell +from ask_shell._internal.live_print_context import LivePrintContext as _LivePrintContext +from ask_shell._internal.live_print_context import get_live_print_context as _get_live_print_context +from ask_shell._internal.live_print_context import live_print_scope as _live_print_scope from ask_shell._internal.rich_live import RemoveLivePart as _RemoveLivePart from ask_shell._internal.rich_live import add_renderable as _add_renderable from ask_shell._internal.rich_live import get_live_console as _get_live_console @@ -11,6 +14,9 @@ disable_interactive_shell = _disable_interactive_shell interactive_shell = _interactive_shell +LivePrintContext = _LivePrintContext +get_live_print_context = _get_live_print_context +live_print_scope = _live_print_scope RemoveLivePart = _RemoveLivePart add_renderable = _add_renderable get_live_console = _get_live_console diff --git a/docs/console/index.md b/docs/console/index.md index a72b35b..e7cc001 100644 --- a/docs/console/index.md +++ b/docs/console/index.md @@ -4,12 +4,15 @@ +- [`LivePrintContext`](#liveprintcontext_def) - [`RemoveLivePart`](#removelivepart_def) - [`add_renderable`](#add_renderable_def) - [`configure_logging`](#configure_logging_def) - [`disable_interactive_shell`](#disable_interactive_shell_def) - [`get_live_console`](#get_live_console_def) +- [`get_live_print_context`](#get_live_print_context_def) - [`interactive_shell`](#interactive_shell_def) +- [`live_print_scope`](#live_print_scope_def) - [`log_to_live`](#log_to_live_def) - [`new_task`](#new_task_def) - [`print_to_live`](#print_to_live_def) @@ -23,7 +26,7 @@ ### class: `RemoveLivePart` -- [source](../../ask_shell/_internal/rich_live.py#L157) +- [source](../../ask_shell/_internal/rich_live.py#L158) > **Since:** 0.3.0 ```python @@ -41,7 +44,7 @@ class RemoveLivePart: ### function: `add_renderable` -- [source](../../ask_shell/_internal/rich_live.py#L161) +- [source](../../ask_shell/_internal/rich_live.py#L162) - [Example: Mount a dynamic Rich renderable in ask-shell's live region with add_renderable, plus an apply-live demo with CI heartbeats](../examples/console/add_renderable.md) > **Since:** 0.3.0 @@ -78,7 +81,7 @@ def configure_logging(app: Typer, *, settings: AskShellSettings | None = None, a ### function: `get_live_console` -- [source](../../ask_shell/_internal/rich_live.py#L179) +- [source](../../ask_shell/_internal/rich_live.py#L180) > **Since:** 0.3.0 ```python @@ -114,7 +117,7 @@ def interactive_shell() -> bool: ### function: `log_to_live` -- [source](../../ask_shell/_internal/rich_live.py#L219) +- [source](../../ask_shell/_internal/rich_live.py#L234) > **Since:** 0.3.0 ```python @@ -168,7 +171,7 @@ class new_task: ### function: `print_to_live` -- [source](../../ask_shell/_internal/rich_live.py#L183) +- [source](../../ask_shell/_internal/rich_live.py#L195) > **Since:** 0.3.0 ```python @@ -201,4 +204,65 @@ Force non-interactive mode for the remainder of this process. | Version | Change | |---------|--------| | 0.9.0 | Made public | - \ No newline at end of file + + + + +### class: `LivePrintContext` +- [source](../../ask_shell/_internal/live_print_context.py#L9) +> **Since:** unreleased + +```python +class LivePrintContext: + prefix: str = '' + suppress: bool = False +``` + +| Field | Type | Default | Since | +|---|---|---|---| +| prefix | `str` | `''` | unreleased | +| suppress | `bool` | `False` | unreleased | + +### Changes + +| Version | Change | +|---------|--------| +| unreleased | Made public | + + + + +### function: `get_live_print_context` +- [source](../../ask_shell/_internal/live_print_context.py#L18) +> **Since:** unreleased + +```python +def get_live_print_context() -> LivePrintContext | None: + ... +``` + +### Changes + +| Version | Change | +|---------|--------| +| unreleased | Made public | + + + + +### function: `live_print_scope` +- [source](../../ask_shell/_internal/live_print_context.py#L22) +- [Example: Prefix or suppress live-console scroll lines per scope, including across run_pool workers](../examples/console/live_print_scope.md) +> **Since:** unreleased + +```python +def live_print_scope(*, prefix: str = '', suppress: bool = False) -> Iterator[None]: + ... +``` + +### Changes + +| Version | Change | +|---------|--------| +| unreleased | Made public | + \ No newline at end of file diff --git a/docs/examples/console/live_print_scope.md b/docs/examples/console/live_print_scope.md new file mode 100644 index 0000000..84d7467 --- /dev/null +++ b/docs/examples/console/live_print_scope.md @@ -0,0 +1,76 @@ + +# live_print_scope + +`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. + +## Read the active scope + +```python +from ask_shell.console import get_live_print_context, live_print_scope + +with live_print_scope(prefix="[prod] "): + ctx = get_live_print_context() + print(repr(ctx.prefix if ctx else "")) + #> '[prod] ' +``` + +## Suppress live output + +```python +from ask_shell.console import get_live_print_context, live_print_scope + +with live_print_scope(suppress=True): + ctx = get_live_print_context() + print(ctx.suppress if ctx else False) + #> True +``` + +Nested scopes restore the outer context when the inner block exits. + +```python +from ask_shell.console import get_live_print_context, live_print_scope + +with live_print_scope(prefix="outer "): + outer = get_live_print_context() + with live_print_scope(prefix="inner "): + inner = get_live_print_context() + after = get_live_print_context() +print(f"{outer.prefix}|{inner.prefix}|{after.prefix}") +#> outer |inner |outer +``` + +## Context at run_pool submit time + +Set `live_print_scope` in the thread that calls `submit`. Each worker sees the scope that was active for its submit call. + +```python +from threading import Barrier + +from ask_shell.console import get_live_print_context, live_print_scope +from ask_shell.shell import run_pool + +results: list[str] = [] +barrier = Barrier(2) + + +def worker() -> None: + barrier.wait() + ctx = get_live_print_context() + results.append(ctx.prefix if ctx else "") + + +with run_pool("demo", total=2, pool_thread_count=2, max_concurrent_submits=2) as pool: + with live_print_scope(prefix="a"): + f_a = pool.submit(worker) + with live_print_scope(prefix="b"): + f_b = pool.submit(worker) + f_a.result() + f_b.result() + +print(sorted(results)) +#> ['a', 'b'] +``` diff --git a/docs/shell/run_pool.md b/docs/shell/run_pool.md index 6ed08c3..111c468 100644 --- a/docs/shell/run_pool.md +++ b/docs/shell/run_pool.md @@ -2,7 +2,7 @@ ## class: run_pool -- [source](../../ask_shell/_internal/run_pool.py#L28) +- [source](../../ask_shell/_internal/run_pool.py#L29) > **Since:** 0.3.0 ```python diff --git a/mkdocs.yml b/mkdocs.yml index ef7b874..8c691d8 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -26,6 +26,7 @@ nav: - Overview: console/index.md - Examples: - add_renderable: examples/console/add_renderable.md + - live_print_scope: examples/console/live_print_scope.md - shell: - Overview: shell/index.md - ShellConfig: shell/shellconfig.md diff --git a/pyproject.toml b/pyproject.toml index da7a5e5..e9d8c4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ examples_include = [ [tool.pkg-ext.groups.console] examples_include = [ "add_renderable", + "live_print_scope", ]