From a4823711562960f94dd50e3ce521eb6355655352 Mon Sep 17 00:00:00 2001 From: "Axel H." Date: Sat, 25 Jul 2026 20:41:49 +0200 Subject: [PATCH 1/2] feat(cli): Rich-based output, prints and debug-inventory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route all human-facing CLI output through Rich, replacing Click as the output backend (Click is now removed entirely). The core stays decoupled via the pluggable `pyinfra.api.output` layer, now wired to Rich implementations. - New `pyinfra_cli.console`: shared Rich stderr console (+ stdout console for `--json`), `format_text` (click.style-compatible signature) and `echo`. - `pyinfra.api.output`: add `get_console`/`set_console`. - `main.py` wires `set_formatter(format_text)` / `set_echo(echo)`. - `prints.py`/`log.py`/`util.py`/`virtualenv.py`: Click echo/style → Rich. - `debug-inventory` renders a Rich table with per-host `key: value` data (nested values as JSON, type-coloured to match the JSON highlighter). - Unexpected errors render a Rich traceback (suppressing framework frames); the internal-error path still writes `pyinfra-debug.log`. - `--help` uses Rich formatting with a colourised usage line and a syntax-highlighted examples epilogue; the confirm prompt uses `rich.prompt`. - Core spinner switches to a shared Rich progress instance. - pyproject: remove `click`, add `rich`; keep `cyclopts`. This is the output layer only; the live progress tree follows separately. --- pyproject.toml | 3 +- src/pyinfra/api/output.py | 51 ++++++-- src/pyinfra_cli/cli.py | 118 ++++++++++++------- src/pyinfra_cli/console.py | 70 +++++++++++ src/pyinfra_cli/exceptions.py | 133 +++++++++++---------- src/pyinfra_cli/log.py | 17 +-- src/pyinfra_cli/main.py | 18 ++- src/pyinfra_cli/prints.py | 189 +++++++++++++++++++++++------- src/pyinfra_cli/util.py | 4 +- src/pyinfra_cli/virtualenv.py | 6 +- tests/test_cli/test_cli.py | 53 +++++++++ tests/test_cli/test_cli_prints.py | 87 ++++++++++++++ tests/test_cli/util.py | 34 ++++-- uv.lock | 8 +- 14 files changed, 596 insertions(+), 195 deletions(-) create mode 100644 src/pyinfra_cli/console.py create mode 100644 tests/test_cli/test_cli_prints.py diff --git a/pyproject.toml b/pyproject.toml index b502922bc..8bbff38c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,8 +14,8 @@ requires-python = ">=3.10,<4.0" dependencies = [ "gevent>=1.5", "paramiko>=2.11,<6", # 2.11 (2022) adds Transport.open_channel(timeout=...) for ProxyJump timeout (#971) - "click>2", "cyclopts>=4,<5", + "rich>=13", "jinja2>3,<4", "python-dateutil>2,<3", "typeguard>=4,<5", @@ -50,7 +50,6 @@ repository = "https://github.com/pyinfra-dev/pyinfra" [dependency-groups] test = [ - "click>=8.2", "pytest>=8.3.5,<9", "freezegun>=1.5.5", "coverage>=7.7.1,<8", diff --git a/src/pyinfra/api/output.py b/src/pyinfra/api/output.py index d7a1bcb68..5f4787edb 100644 --- a/src/pyinfra/api/output.py +++ b/src/pyinfra/api/output.py @@ -3,14 +3,50 @@ Provides ``format_text`` and ``echo`` functions that default to plain-text no-ops, allowing the API layer to work without any CLI dependency. The CLI -layer replaces them at startup via ``set_formatter`` and ``set_echo``. +layer replaces them at startup via ``set_formatter`` and ``set_echo`` (wiring +in Rich-backed implementations). """ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any from collections.abc import Callable +if TYPE_CHECKING: + from rich.console import Console + +_console: Console | None = None + + +def get_console() -> Console: + """ + Return the shared human-facing (stderr) Rich console. + + Created lazily so importing the API layer doesn't require Rich to be + configured. The CLI may replace it via :func:`set_console` to share a + single console between logging, tables and the progress spinner. + """ + global _console + if _console is None: + from rich.console import Console + + # markup/emoji disabled: host print prefixes like ``[@fake/host]`` and + # arbitrary command output must not be interpreted as Rich markup. + _console = Console( + stderr=True, + highlight=False, + soft_wrap=True, + markup=False, + emoji=False, + ) + return _console + + +def set_console(console: Console) -> None: + """Replace the shared human-facing console.""" + global _console + _console = console + # Default formatter: identity function (returns plain text, ignores styling kwargs). def _default_format_text(text: str, *args: Any, **kwargs: Any) -> str: @@ -29,25 +65,26 @@ def _default_echo(message: Any = None, **kwargs: Any) -> None: def format_text(text: str, *args: Any, **kwargs: Any) -> str: """Format text with optional styling (color, bold, etc.). - Mirrors ``click.style`` signature. Accepts positional ``fg`` argument - for compatibility with ``click.style("text", "red")``. + Historically mirrored ``click.style``: accepts a positional foreground + color (e.g. ``format_text("text", "red")``) and ``bold=`` keyword. The CLI + installs a Rich-backed implementation preserving this signature. """ return _format_text(text, *args, **kwargs) def echo(message: Any = None, **kwargs: Any) -> None: - """Echo a message. Mirrors ``click.echo`` signature.""" + """Echo a message. Supports ``err=True`` to write to stderr.""" return _echo(message, **kwargs) def set_formatter(func: Callable[..., str]) -> None: - """Replace the default formatter (e.g. with ``click.style``).""" + """Replace the default formatter (e.g. with a Rich-backed styler).""" global _format_text _format_text = func def set_echo(func: Callable[..., None]) -> None: - """Replace the default echo function (e.g. with ``click.echo``).""" + """Replace the default echo function (e.g. with a Rich-backed echo).""" global _echo _echo = func diff --git a/src/pyinfra_cli/cli.py b/src/pyinfra_cli/cli.py index 7789adddc..3194b6c63 100644 --- a/src/pyinfra_cli/cli.py +++ b/src/pyinfra_cli/cli.py @@ -9,8 +9,8 @@ from os import chdir as os_chdir, environ, getcwd from pathlib import Path -import click from cyclopts import App, Group, Parameter +from rich.prompt import Confirm from pyinfra import __version__, logger, state from pyinfra.api import Config, Host, Inventory, State @@ -26,6 +26,7 @@ from pyinfra.api.output import format_text from .commands import get_facts_and_args, get_func_and_args +from .console import console, stdout_console from .exceptions import CliError, UnexpectedExternalError, UnexpectedInternalError, WrappedError from .inventory import make_inventory from .log import setup_logging @@ -65,11 +66,71 @@ def _lenient_bool(type_, tokens) -> bool: raise ValueError(f"invalid boolean value: {tokens[0].value!r}") +_EXAMPLES = """\ +# Run one or more deploys against the inventory +pyinfra INVENTORY deploy_web.py [deploy_db.py]... + +# Run a single operation against the inventory +pyinfra INVENTORY server.user pyinfra home=/home/pyinfra + +# Execute an arbitrary command against the inventory +pyinfra INVENTORY exec -- echo "hello world" + +# Run one or more facts against the inventory +pyinfra INVENTORY fact server.LinuxName [server.Users]... +pyinfra INVENTORY fact files.File path=/path/to/file... + +# Debug the inventory hosts and data +pyinfra INVENTORY debug-inventory""" + + +def _build_examples_epilogue() -> str: + """Render the CLI examples as syntax-highlighted (bash) ANSI text. + + Used as a ``help_format="rich"`` epilogue so the examples show up + colourised on the help page. + """ + from rich.syntax import Syntax + + syntax = Syntax(_EXAMPLES, "bash", background_color="default", word_wrap=True) + with stdout_console.capture() as capture: + stdout_console.print("[bold]Examples:[/bold]\n") + stdout_console.print(syntax) + return capture.get() + + +def _build_usage() -> str: + """Colourised usage line: required args in cyan, optionals dimmed. + + Rendered to an ANSI string because Cyclopts concatenates ``usage`` as a + plain string; the ``Usage:`` label is added by the help formatter. Colour + is only emitted when stdout is a terminal so piped output stays plain. + """ + from rich.text import Text + + line = Text.assemble( + ("pyinfra ", "bold"), + ("[OPTIONS] ", "dim"), + ("INVENTORY ", "bold cyan"), + ("[OPERATIONS...]", "cyan"), + ) + if not stdout_console.is_terminal: + return line.plain + with stdout_console.capture() as capture: + stdout_console.print(line, end="") + return capture.get() + + app = App( name="pyinfra", version=f"pyinfra: v{__version__}", version_flags=["--version"], help_flags=["-h", "--help"], + help_format="rich", + usage=_build_usage(), + console=stdout_console, + error_console=console, + help_epilogue=_build_examples_epilogue(), ) # Enable ``pyinfra --install-completion`` for shell autocompletion. @@ -169,30 +230,14 @@ def cli( """pyinfra manages the state of one or more servers. It can be used for app/service deployment, config management and ad-hoc - command execution. Documentation: docs.pyinfra.com - - INVENTORY is a file (inventory.py), a hostname (host.net) or comma separated - hostnames (host-1.net,host-2.net,@local). + command execution. - Examples: + Documentation: [cyan][link=https://docs.pyinfra.com]docs.pyinfra.com[/link][/cyan] - ``` - # Run one or more deploys against the inventory - pyinfra INVENTORY deploy_web.py [deploy_db.py]... - - # Run a single operation against the inventory - pyinfra INVENTORY server.user pyinfra home=/home/pyinfra - - # Execute an arbitrary command against the inventory - pyinfra INVENTORY exec -- echo "hello world" - - # Run one or more facts against the inventory - pyinfra INVENTORY fact server.LinuxName [server.Users]... - pyinfra INVENTORY fact files.File path=/path/to/file... - - # Debug the inventory hosts and data - pyinfra INVENTORY debug-inventory - ``` + INVENTORY can be: + - a file ([cyan]inventory.py[/cyan]) + - a hostname ([cyan]host.net[/cyan]) + - comma separated hostnames ([cyan]host-1.net,host-2.net,@local[/cyan]) Parameters ---------- @@ -489,13 +534,12 @@ def _main( else: logger.info("--> Detected changes:") print_meta(state) - click.echo( + console.print( """ Detected changes may not include every change pyinfra will execute. Hidden side effects of operations may alter behaviour of future operations, this will be shown in the results. The remote state will always be updated to reflect the state defined by the input operations.""", - err=True, ) # If --debug-facts or --debug-operations, print and exit @@ -536,29 +580,15 @@ def _main( def _do_confirm(msg: str) -> bool: - click.echo(err=True) - click.echo(f" {msg}", err=True) + console.print() + console.print(f" {msg}") warning_count = state.get_warning_counter() if warning_count > 0: - click.secho( + console.print( f" {warning_count} warnings shown during change detection, see above", - fg="yellow", - err=True, + style="yellow", ) - confirm_msg = " Press enter to execute..." - click.echo(confirm_msg, err=True, nl=False) - v = input() - if v: - click.echo(f" Unexpected user input: {v}", err=True) - return False - # Go up, clear the line, go up again - as if the confirmation statement was never here! - click.echo( - "\033[1A{}\033[1A".format("".join(" " for _ in range(len(confirm_msg)))), - err=True, - nl=False, - ) - click.echo(err=True) - return True + return Confirm.ask(" Execute?", console=console, default=True) # Setup diff --git a/src/pyinfra_cli/console.py b/src/pyinfra_cli/console.py new file mode 100644 index 000000000..8fb719eae --- /dev/null +++ b/src/pyinfra_cli/console.py @@ -0,0 +1,70 @@ +""" +Shared Rich consoles and Click-compatible output adapters for the CLI. + +pyinfra keeps all human-facing output on **stderr** and reserves **stdout** for +machine-readable (``--json``) payloads. The core library styles/echoes text +through :mod:`pyinfra.api.output`; here we install Rich-backed implementations. +""" + +from __future__ import annotations + +from typing import Any + +from rich.console import Console +from rich.text import Text + +from pyinfra.api.output import get_console, set_console + +# Human-facing console (logs, tables, prompts, spinner) → stderr. +# Reuse the core shared console so the progress spinner and logging write to the +# same Console instance (avoids Live-region corruption). +console = get_console() +set_console(console) + +# Machine-readable console (``--json`` payloads) → stdout, no styling. +stdout_console = Console(highlight=False, soft_wrap=True, markup=False, emoji=False) + + +def format_text(text: str, fg: str | None = None, *, bold: bool = False, **kwargs: Any) -> str: + """ + Style ``text`` and return a string with embedded ANSI codes. + + Mirrors the legacy ``click.style`` signature (positional foreground color + + ``bold=``) used across the core library, but renders via Rich so styling is + consistent with the rest of the CLI output. Colour names (``red``, + ``green``, ...) are passed straight through to Rich. + """ + style_bits = [] + if fg is not None: + style_bits.append(fg) + if bold: + style_bits.append("bold") + + if not style_bits: + return text + + rich_text = Text(text, style=" ".join(style_bits)) + with console.capture() as capture: + console.print(rich_text, end="") + return capture.get() + + +def echo(message: Any = None, *, err: bool = False, nl: bool = True, **kwargs: Any) -> None: + """ + Print ``message`` to the appropriate console. + + ``err=True`` targets the (default) human stderr console; ``err=False`` + targets stdout. ``nl=False`` suppresses the trailing newline. Text may + contain ANSI escape codes already produced by :func:`format_text`. + """ + target = console if err else stdout_console + end = "\n" if nl else "" + + if message is None: + target.print("", end=end) + return + + if isinstance(message, str): + target.print(Text.from_ansi(message), end=end) + else: + target.print(message, end=end) diff --git a/src/pyinfra_cli/exceptions.py b/src/pyinfra_cli/exceptions.py index 9f2d3205b..c8643554b 100644 --- a/src/pyinfra_cli/exceptions.py +++ b/src/pyinfra_cli/exceptions.py @@ -1,12 +1,13 @@ -import abc import sys from inspect import getframeinfo -from traceback import format_exception, format_tb, walk_tb -from types import TracebackType +from traceback import walk_tb +from types import ModuleType, TracebackType -import click +from rich.console import Console +from rich.traceback import Traceback from typing_extensions import override +import pyinfra from pyinfra import logger from pyinfra.api.exceptions import ( ArgumentTypeError, @@ -16,6 +17,29 @@ ) from pyinfra.api.util import PYINFRA_INSTALL_DIR +from .console import console, format_text + +# Modules whose frames are collapsed in rendered tracebacks so the user's deploy +# code stands out rather than pyinfra/gevent/cyclopts internals. +_TRACEBACK_SUPPRESS: list[str | ModuleType] = ["gevent", "cyclopts", pyinfra] + + +def _rich_traceback(exc: BaseException) -> Traceback: + """Build a Rich ``Traceback`` for a wrapped exception. + + The wrapping ``CliException`` stashes the live traceback on the original + exception as ``_traceback``; fall back to ``__traceback__`` just in case. + """ + tb = getattr(exc, "_traceback", None) or exc.__traceback__ + return Traceback.from_exception( + type(exc), + exc, + tb, + suppress=_TRACEBACK_SUPPRESS, + show_locals=False, + word_wrap=True, + ) + def get_frame_line_from_tb(tb: TracebackType): frame_lines = list(walk_tb(tb)) @@ -27,7 +51,24 @@ def get_frame_line_from_tb(tb: TracebackType): return info -class WrappedError(click.ClickException): +class CliException(Exception): + """Base for pyinfra CLI errors, carrying a user-facing ``message``.""" + + message: str + + def __init__(self, message: str = ""): + self.message = message + super().__init__(message) + + @override + def __str__(self) -> str: + return self.message + + def show(self) -> None: + raise NotImplementedError + + +class WrappedError(CliException): def __init__(self, e: Exception): self.traceback = e.__traceback__ self.exception = e @@ -36,10 +77,10 @@ def __init__(self, e: Exception): message = getattr(e, "message", e.args[0]) if not isinstance(message, str): message = repr(message) - self.message = message + super().__init__(message) @override - def show(self, file=None): + def show(self) -> None: name = "unknown error" if isinstance(self.exception, ConnectorDataTypeError): @@ -59,45 +100,31 @@ def show(self, file=None): name = f"{name} in {info.filename} line {info.lineno}" logger.warning( - f"--> {click.style(name, 'red', bold=True)}: {self}", + f"--> {format_text(name, 'red', bold=True)}: {self}", ) -class CliError(click.ClickException): +class CliError(CliException): @override - def show(self, file=None): + def show(self) -> None: logger.warning( - f"--> {click.style('pyinfra error', 'red', bold=True)}: {self}", + f"--> {format_text('pyinfra error', 'red', bold=True)}: {self}", ) -class UnexpectedMixin(abc.ABC): - exception: Exception - traceback: TracebackType - - def get_traceback_lines(self): - traceback = getattr(self.exception, "_traceback") - return format_tb(traceback) - - def get_traceback(self): - return "".join(self.get_traceback_lines()) - - def get_exception(self): - return "".join(format_exception(self.exception.__class__, self.exception, None)) - - -class UnexpectedExternalError(click.ClickException, UnexpectedMixin): +class UnexpectedExternalError(CliException): def __init__(self, e, filename): _, _, traceback = sys.exc_info() e._traceback = traceback self.exception = e self.filename = filename + super().__init__(str(e)) @override - def show(self, file=None): + def show(self) -> None: logger.warning( "--> {}:\n".format( - click.style( + format_text( f"An exception occurred in: {self.filename}", "red", bold=True, @@ -105,56 +132,42 @@ def show(self, file=None): ), ) - click.echo("Traceback (most recent call last):", err=True) - click.echo(self.get_traceback(), err=True, nl=False) - click.echo(self.get_exception(), err=True) + console.print(_rich_traceback(self.exception)) -class UnexpectedInternalError(click.ClickException, UnexpectedMixin): +class UnexpectedInternalError(CliException): def __init__(self, e): _, _, traceback = sys.exc_info() e._traceback = traceback self.exception = e + super().__init__(str(e)) @override - def show(self, file=None): - click.echo( + def show(self) -> None: + console.print( "--> {}:\n".format( - click.style( + format_text( "An internal exception occurred", "red", bold=True, ), ), - err=True, ) - traceback_lines = self.get_traceback_lines() - traceback = self.get_traceback() - - # Syntax errors contain the filename/line/etc, but other exceptions - # don't, so print the *last* call to stderr. - if not isinstance(self.exception, SyntaxError): - sys.stderr.write(traceback_lines[-1]) - - exception = self.get_exception() - click.echo(exception, err=True) + traceback = _rich_traceback(self.exception) + console.print(traceback) + # Persist an uncoloured copy of the same traceback for bug reports. with open("pyinfra-debug.log", "w", encoding="utf-8") as f: - f.write(traceback) - f.write(exception) + file_console = Console(file=f, width=100, force_terminal=False, no_color=True) + file_console.print(traceback) - logger.debug(traceback) - logger.debug(exception) + logger.debug(str(self.exception)) - click.echo( - f"--> The full traceback has been written to {click.style('pyinfra-debug.log', bold=True)}", - err=True, + console.print( + f"--> The full traceback has been written to {format_text('pyinfra-debug.log', bold=True)}", ) - click.echo( - ( - "--> If this is unexpected please consider submitting a bug report " - "on GitHub, for more information run `pyinfra --support`." - ), - err=True, + console.print( + "--> If this is unexpected please consider submitting a bug report " + "on GitHub, for more information run `pyinfra --support`." ) diff --git a/src/pyinfra_cli/log.py b/src/pyinfra_cli/log.py index ece7e0c5b..a64a48e41 100644 --- a/src/pyinfra_cli/log.py +++ b/src/pyinfra_cli/log.py @@ -1,18 +1,21 @@ import logging -import click +from rich.text import Text from typing_extensions import override from pyinfra import logger, state from pyinfra.context import ctx_state +from .console import console, format_text + class LogHandler(logging.Handler): @override def emit(self, record): try: message = self.format(record) - click.echo(message, err=True) + # ``message`` may already contain ANSI escape codes (from format_text). + console.print(Text.from_ansi(message)) except Exception: self.handleError(record) @@ -21,10 +24,10 @@ class LogFormatter(logging.Formatter): previous_was_header = True level_to_format = { - logging.DEBUG: lambda s: click.style(s, "green"), - logging.WARNING: lambda s: click.style(s, "yellow"), - logging.ERROR: lambda s: click.style(s, "red"), - logging.CRITICAL: lambda s: click.style(s, "red", bold=True), + logging.DEBUG: lambda s: format_text(s, "green"), + logging.WARNING: lambda s: format_text(s, "yellow"), + logging.ERROR: lambda s: format_text(s, "red"), + logging.CRITICAL: lambda s: format_text(s, "red", bold=True), } @override @@ -50,7 +53,7 @@ def format(self, record): if "-->" in message: if not self.previous_was_header: - click.echo(err=True) + console.print() else: message = f" {message}" diff --git a/src/pyinfra_cli/main.py b/src/pyinfra_cli/main.py index 289fbb228..5c3ee3e04 100644 --- a/src/pyinfra_cli/main.py +++ b/src/pyinfra_cli/main.py @@ -1,37 +1,35 @@ import signal import sys -import click import gevent import pyinfra from pyinfra.api.output import set_echo, set_formatter from .cli import app +from .console import console, echo, format_text +from .exceptions import CliException def main(): # Set CLI mode pyinfra.is_cli = True - # Wire click's styling/echo into the API output layer - set_formatter(click.style) - set_echo(click.echo) + # Wire Rich-backed styling/echo into the API output layer + set_formatter(format_text) + set_echo(echo) # Don't write out deploy.pyc/config.pyc etc sys.dont_write_bytecode = True sys.path.append(".") - # Shut it click - click.disable_unicode_literals_warning = True # type: ignore - # Force line buffering sys.stdout.reconfigure(line_buffering=True) # type: ignore sys.stderr.reconfigure(line_buffering=True) # type: ignore def _handle_interrupt(signum, frame): - click.echo("Exiting upon user request!") + console.print("Exiting upon user request!") sys.exit(0) try: @@ -45,6 +43,6 @@ def _handle_interrupt(signum, frame): try: app() - except click.ClickException as e: + except CliException as e: e.show() - sys.exit(e.exit_code) + sys.exit(1) diff --git a/src/pyinfra_cli/prints.py b/src/pyinfra_cli/prints.py index 625017aef..f85cfa4c9 100644 --- a/src/pyinfra_cli/prints.py +++ b/src/pyinfra_cli/prints.py @@ -4,14 +4,20 @@ import platform import re import sys -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from collections.abc import Callable, Iterator -import click +from rich.console import Group +from rich.json import JSON +from rich.padding import Padding +from rich.table import Table +from rich.text import Text from pyinfra import __version__, logger from pyinfra.api.host import Host +from pyinfra.api.output import format_text +from .console import console, stdout_console from .util import json_encode if TYPE_CHECKING: @@ -53,8 +59,50 @@ def jsonify(data, *args, **kwargs): return json.dumps(data, *args, **kwargs) +def _safe_encode(obj: Any) -> Any: + """``json_encode`` fallback that never raises (for values). + + Used for the human ``debug-inventory`` rendering, where a value that is + neither natively JSON-serialisable nor handled by ``json_encode`` (e.g. a + compiled ``re.Pattern``) should degrade to its ``str()`` rather than + aborting the whole command. The ``--json`` path keeps using the strict + ``json_encode`` so machine output stays valid JSON. + """ + try: + return json_encode(obj) + except TypeError: + return str(obj) + + +def _json_safe_keys(value: Any) -> Any: + """Recursively coerce non-primitive mapping keys to ``str``. + + ``json.dumps`` rejects dict keys that are not ``str``/``int``/``float``/ + ``bool``/``None`` *before* the ``default`` hook runs, so a ``re.Pattern`` + used as a ``fake_responses`` matcher key would still raise. This makes the + human ``debug-inventory`` rendering robust against such keys. + """ + if isinstance(value, dict): + return { + (key if isinstance(key, (str, int, float, bool)) or key is None else str(key)): ( + _json_safe_keys(val) + ) + for key, val in value.items() + } + if isinstance(value, (list, tuple)): + return [_json_safe_keys(item) for item in value] + return value + + def print_json(payload) -> None: - click.echo(jsonify(payload, default=json_encode)) + json_str = jsonify(payload, default=json_encode) + + # When stdout is a real terminal, pretty-print + syntax-highlight the JSON. + # When piped/redirected, emit plain JSON so it stays machine-parseable. + if stdout_console.is_terminal: + stdout_console.print(JSON(json_str)) + else: + print(json_str) def _host_to_dict(host: Host) -> dict: @@ -192,23 +240,22 @@ def print_run_json(state: State, dry: bool) -> None: def print_state_operations(state: State): state_ops = {host: ops for host, ops in state.ops.items() if state.is_host_in_limit(host)} - click.echo(err=True) - click.echo("--> Operations:", err=True) - click.echo(jsonify(state_ops, indent=4, default=json_encode), err=True) - click.echo(err=True) - click.echo("--> Operation meta:", err=True) - click.echo(jsonify(state.op_meta, indent=4, default=json_encode), err=True) + console.print() + console.print("--> Operations:") + console.print(jsonify(state_ops, indent=4, default=json_encode)) + console.print() + console.print("--> Operation meta:") + console.print(jsonify(state.op_meta, indent=4, default=json_encode)) - click.echo(err=True) - click.echo("--> Operation order:", err=True) - click.echo(err=True) + console.print() + console.print("--> Operation order:") + console.print() for op_hash in state.get_op_order(): meta = state.op_meta[op_hash] hosts = set(host for host, operations in state.ops.items() if op_hash in operations) - click.echo( + console.print( f" {op_hash} (names={meta.names}, hosts={hosts})", - err=True, ) @@ -222,9 +269,8 @@ def print_groups_by_comparison(print_items, comparator=lambda item: item[0]): items.append(name) else: - click.echo( - f" {', '.join(click.style(name, bold=True) for name in items)}", - err=True, + console.print( + f" {', '.join(format_text(name, bold=True) for name in items)}", ) items = [name] @@ -232,31 +278,94 @@ def print_groups_by_comparison(print_items, comparator=lambda item: item[0]): last_name = name if items: - click.echo( - f" {', '.join(click.style(name, bold=True) for name in items)}", - err=True, + console.print( + f" {', '.join(format_text(name, bold=True) for name in items)}", ) def print_fact(fact_data): - click.echo(jsonify(fact_data, indent=4, default=json_encode), err=True) + console.print(jsonify(fact_data, indent=4, default=json_encode)) + + +def _scalar_style(value: Any) -> str: + """Rich style for a scalar, matching the JSON highlighter's type colours. + + Non-JSON scalars (datetime, Path, ``re.Pattern``, arbitrary objects) render + unstyled, since they are shown via ``str()`` rather than as JSON values. + """ + # NOTE: bool is a subclass of int, so it must be checked first. + if isinstance(value, bool): + return "json.bool_true" if value else "json.bool_false" + if value is None: + return "json.null" + if isinstance(value, (int, float)): + return "json.number" + if isinstance(value, str): + return "json.str" + return "" + + +def _format_host_data(data: dict) -> Group: + """Render host data as one ``key: value`` line per top-level key. + + Scalars are shown inline; nested ``dict``/``list``/``tuple`` values are + rendered as indented JSON (syntax-highlighted). Any other value (datetime, + Path, ``re.Pattern``, arbitrary objects) falls back to ``str()`` so the + display never fails on non-JSON-serialisable data. Insertion order is + preserved. + """ + if not data: + return Group(Text("(no data)", style="dim")) + + lines: list[Any] = [] + for key, value in data.items(): + label = Text(f"{key}: ", style="bold blue") + if isinstance(value, (dict, list, tuple)): + # Nested structures: header line + indented JSON below it. + lines.append(Text.assemble(label)) + value_json = jsonify(_json_safe_keys(value), indent=2, default=_safe_encode) + lines.append(Padding(JSON(value_json), (0, 0, 0, 2))) + else: + # Scalars inline, coloured to match Rich's JSON highlighter (booleans + # green/red, numbers cyan, null magenta, strings green). Other values + # (datetime, Path, re.Pattern, arbitrary objects) fall back to an + # unstyled str() so the display never fails on non-JSON data. + lines.append(Text.assemble(label, (str(value), _scalar_style(value)))) + + return Group(*lines) def print_inventory(state: State): + table = Table( + title="Inventory", + title_style="bold", + header_style="bold", + expand=True, + leading=1, + ) + # Only the data column flexes; host/groups stay as narrow as their content. + table.add_column("Host", style="cyan", no_wrap=True, ratio=None) + table.add_column("Groups", style="green", no_wrap=True, ratio=None) + table.add_column("Data", ratio=1) + for host in state.inventory: - click.echo(err=True) - click.echo(host.print_prefix, err=True) - click.echo(f"--> Groups: {', '.join(host.groups)}", err=True) - click.echo("--> Data:", err=True) - click.echo(jsonify(host.data, indent=4, default=json_encode), err=True) + # A host may appear in the same group more than once (e.g. connector + + # inventory group); de-duplicate for display while preserving order. + groups = list(dict.fromkeys(host.groups)) + table.add_row( + host.name, + "\n".join(groups), + _format_host_data(host.data.dict()), + ) + + console.print(table) def print_facts(facts): for name, data in facts.items(): - click.echo(err=True) - click.echo( - f"--> Fact data for: {click.style(name, bold=True)}", - err=True, + console.print() + console.print( + f"--> Fact data for: {format_text(name, bold=True)}", ) print_fact(data) @@ -266,7 +375,7 @@ def print_support_info() -> None: from packaging.requirements import Requirement - click.echo( + console.print( """ If you are having issues with pyinfra or wish to make feature requests, please check out the GitHub issues at https://github.com/Fizzadar/pyinfra/issues . @@ -274,11 +383,11 @@ def print_support_info() -> None: """, ) - click.echo(f" System: {platform.system()}", err=True) - click.echo(f" Platform: {platform.platform()}", err=True) - click.echo(f" Release: {platform.uname()[2]}", err=True) - click.echo(f" Machine: {platform.uname()[4]}", err=True) - click.echo(f" pyinfra: v{__version__}", err=True) + console.print(f" System: {platform.system()}") + console.print(f" Platform: {platform.platform()}") + console.print(f" Release: {platform.uname()[2]}") + console.print(f" Machine: {platform.uname()[4]}") + console.print(f" pyinfra: v{__version__}") seen_reqs: set[str] = set() for requirement_string in sorted(requires("pyinfra") or []): @@ -287,18 +396,16 @@ def print_support_info() -> None: continue seen_reqs.add(requirement.name) try: - click.echo( + console.print( f" {requirement.name}: v{version(requirement.name)}", - err=True, ) except PackageNotFoundError: # package not installed in this environment continue - click.echo(f" Executable: {sys.argv[0]}", err=True) - click.echo( + console.print(f" Executable: {sys.argv[0]}") + console.print( f" Python: {platform.python_version()} ({platform.python_implementation()}, {platform.python_compiler()})", - err=True, ) diff --git a/src/pyinfra_cli/util.py b/src/pyinfra_cli/util.py index 5bea669f4..cdaefd306 100644 --- a/src/pyinfra_cli/util.py +++ b/src/pyinfra_cli/util.py @@ -9,7 +9,6 @@ from types import CodeType, FunctionType, ModuleType from collections.abc import Callable -import click import gevent from pyinfra import logger, state @@ -17,6 +16,7 @@ from pyinfra.api.exceptions import PyinfraError from pyinfra.api.host import HostData from pyinfra.api.operation import OperationMeta +from pyinfra.api.output import format_text from pyinfra.api.state import ( State, StateHostMeta, @@ -221,7 +221,7 @@ def load_file(local_host): with ctx_host.use(local_host): callback() logger.info( - f"{local_host.print_prefix}{click.style('Ready:', 'green')} {click.style(name, bold=True)}", + f"{local_host.print_prefix}{format_text('Ready:', 'green')} {format_text(name, bold=True)}", ) except Exception as e: return e diff --git a/src/pyinfra_cli/virtualenv.py b/src/pyinfra_cli/virtualenv.py index ed59916bd..eed0fa1bb 100644 --- a/src/pyinfra_cli/virtualenv.py +++ b/src/pyinfra_cli/virtualenv.py @@ -2,10 +2,10 @@ import sys from pathlib import Path -import click - from pyinfra import logger +from .console import console + def init_virtualenv() -> None: """ @@ -62,7 +62,7 @@ def init_virtualenv() -> None: " If you encounter problems, please install pyinfra inside the virtualenv." ), ) - click.echo(err=True) + console.print() if sys.platform == "win32": virtual_env = str(Path(os.environ["VIRTUAL_ENV"]) / "Lib" / "site-packages") diff --git a/tests/test_cli/test_cli.py b/tests/test_cli/test_cli.py index e4ee19701..2349d5e3f 100644 --- a/tests/test_cli/test_cli.py +++ b/tests/test_cli/test_cli.py @@ -17,6 +17,59 @@ def test_print_help(self): result = run_cli("--help") assert result.exit_code == 0, result.stderr + def test_support_standalone(self): + # `pyinfra --support` must work without INVENTORY/OPERATIONS: the crash + # handler tells users to run exactly this. + result = run_cli("--support") + assert result.exit_code == 0, result.stderr + assert "pyinfra: v" in result.stderr + + +class TestCliYesEnvVar(TestCase): + def _parse_yes(self, value): + import os + + from pyinfra_cli.cli import app + + os.environ["PYINFRA_YES"] = value + try: + _, bound, _ = app.parse_args(["inv.py", "server.shell", "x"], exit_on_error=False) + return bound.arguments["yes"] + finally: + del os.environ["PYINFRA_YES"] + + def test_empty_is_false(self): + assert self._parse_yes("") is False + + def test_whitespace_is_false(self): + assert self._parse_yes(" ") is False + + def test_on_off(self): + assert self._parse_yes("on") is True + assert self._parse_yes("off") is False + + def test_numeric(self): + assert self._parse_yes("1") is True + assert self._parse_yes("0") is False + + def test_true_false_any_case(self): + assert self._parse_yes("true") is True + assert self._parse_yes("False") is False + + def test_invalid_value_errors(self): + import os + + from cyclopts import CycloptsError + + from pyinfra_cli.cli import app + + os.environ["PYINFRA_YES"] = "junk" + try: + with self.assertRaises(CycloptsError): + app.parse_args(["inv.py", "x"], exit_on_error=False, print_error=False) + finally: + del os.environ["PYINFRA_YES"] + class TestOperationCli(PatchSSHTestCase): def test_invalid_operation_module(self): diff --git a/tests/test_cli/test_cli_prints.py b/tests/test_cli/test_cli_prints.py new file mode 100644 index 000000000..9be48e843 --- /dev/null +++ b/tests/test_cli/test_cli_prints.py @@ -0,0 +1,87 @@ +import re +from datetime import datetime +from pathlib import PurePosixPath +from unittest import TestCase +from unittest.mock import patch + +from rich.console import Console + +from pyinfra.api import Config, State + +import pyinfra_cli.prints as prints_module +from pyinfra_cli.console import console +from pyinfra_cli.prints import _format_host_data, _scalar_style, print_inventory + +from ..util import make_inventory + + +def _render_inventory(host_data: dict) -> str: + inventory = make_inventory(hosts=(("somehost", host_data),)) + state = State(inventory, Config()) + # Use a wide, fixed-width console so the table is never truncated (the shared + # console's width varies by platform/terminal, which would clip cell content). + wide_console = Console(width=200, force_terminal=False, highlight=False) + with patch.object(prints_module, "console", wide_console): + with wide_console.capture() as capture: + print_inventory(state) + return capture.get() + + +class TestPrintInventory(TestCase): + def test_scalars_render_as_flat_key_value_lines(self): + output = _render_inventory({"role": "web", "port": 80, "enabled": True}) + + assert "role: web" in output + assert "port: 80" in output + assert "enabled: True" in output + # Scalars must NOT be dumped as JSON (no quoted keys/values). + assert '"role"' not in output + assert '"web"' not in output + + def test_nested_values_render_as_json(self): + output = _render_inventory({"tags": ["a", "b"], "meta": {"cpu": 4}}) + + # Header line for the key, then indented JSON for the value. + assert "tags:" in output + assert '"a"' in output and '"b"' in output + assert "meta:" in output + assert '"cpu": 4' in output + + def test_non_json_scalars_render_via_str(self): + created = datetime(2021, 8, 14, 10, 30) + # PurePosixPath keeps str() stable across platforms (WindowsPath would + # render with backslashes). + path = PurePosixPath("/opt/app") + output = _render_inventory({"created": created, "path": path}) + + assert f"created: {created}" in output + assert f"path: {path}" in output + + def test_re_pattern_does_not_crash(self): + # Regression: a compiled regex (as a nested dict key) previously crashed + # the whole `debug-inventory` command trying to JSON-encode it. + output = _render_inventory( + {"fake_responses": {re.compile(r"^pip"): {"success": False}}}, + ) + + assert "fake_responses:" in output + assert "success" in output + + def test_empty_data(self): + # `make_inventory` always injects some data, so exercise the helper + # directly for the empty case. + with console.capture() as capture: + console.print(_format_host_data({})) + assert "(no data)" in capture.get() + + def test_scalar_styling_matches_json_highlighter(self): + # Scalars are coloured by type to match Rich's JSON highlighter. + assert _scalar_style(True) == "json.bool_true" + assert _scalar_style(False) == "json.bool_false" + assert _scalar_style(None) == "json.null" + assert _scalar_style(80) == "json.number" + assert _scalar_style(1.5) == "json.number" + assert _scalar_style("web") == "json.str" + # Non-JSON scalars render unstyled (shown via str()). + assert _scalar_style(datetime(2021, 8, 14)) == "" + assert _scalar_style(PurePosixPath("/opt/app")) == "" diff --git a/tests/test_cli/util.py b/tests/test_cli/util.py index e27c500e0..94d9cf4b2 100644 --- a/tests/test_cli/util.py +++ b/tests/test_cli/util.py @@ -2,10 +2,10 @@ from io import StringIO from os import chdir, getcwd -import click - import pyinfra +import pyinfra_cli.console as cli_console from pyinfra_cli.cli import app +from pyinfra_cli.exceptions import CliException class CliResult: @@ -26,26 +26,34 @@ def run_cli(*arguments): stdout_buffer = StringIO() stderr_buffer = StringIO() + # The whole CLI (cli/prints/log/virtualenv/exceptions/progress) shares the + # single console instance from pyinfra_cli.console, so redirecting its file + # captures all human output. Machine-readable (--json) output goes to real + # stdout via print(), captured with redirect_stdout below. + console = cli_console.console + stdout_console = cli_console.stdout_console + original_console_file = console.file + original_stdout_console_file = stdout_console.file + console.file = stderr_buffer + stdout_console.file = stdout_buffer + exit_code = 0 exception = None try: - with ( - contextlib.redirect_stdout(stdout_buffer), - contextlib.redirect_stderr(stderr_buffer), - ): - try: - app(list(arguments), exit_on_error=False) - except click.ClickException as e: - exception = e - e.show() - exit_code = e.exit_code + with contextlib.redirect_stdout(stdout_buffer): + app(list(arguments), exit_on_error=False) except SystemExit as e: exit_code = e.code if isinstance(e.code, int) else (0 if e.code is None else 1) - except BaseException as e: # surface any error to the test as .exception + except CliException as e: + exception = e + exit_code = 1 + except BaseException as e: # noqa: B036 - surface any error to the test as .exception exception = e exit_code = 1 finally: + console.file = original_console_file + stdout_console.file = original_stdout_console_file pyinfra.is_cli = False chdir(cwd) diff --git a/uv.lock b/uv.lock index 6cd73627f..221a6706d 100644 --- a/uv.lock +++ b/uv.lock @@ -1389,7 +1389,6 @@ wheels = [ name = "pyinfra" source = { editable = "." } dependencies = [ - { name = "click" }, { name = "cyclopts" }, { name = "distro" }, { name = "gevent" }, @@ -1398,6 +1397,7 @@ dependencies = [ { name = "paramiko" }, { name = "pydantic" }, { name = "python-dateutil" }, + { name = "rich" }, { name = "typeguard" }, { name = "types-paramiko" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, @@ -1405,7 +1405,6 @@ dependencies = [ [package.dev-dependencies] dev = [ - { name = "click" }, { name = "coverage" }, { name = "freezegun" }, { name = "ipdb" }, @@ -1433,7 +1432,6 @@ docs = [ { name = "zensical" }, ] test = [ - { name = "click" }, { name = "coverage" }, { name = "freezegun" }, { name = "mypy" }, @@ -1451,7 +1449,6 @@ test = [ [package.metadata] requires-dist = [ - { name = "click", specifier = ">2" }, { name = "cyclopts", specifier = ">=4,<5" }, { name = "distro", specifier = ">=1.6,<2" }, { name = "gevent", specifier = ">=1.5" }, @@ -1460,6 +1457,7 @@ requires-dist = [ { name = "paramiko", specifier = ">=2.11,<6" }, { name = "pydantic", specifier = ">=2.11,<3" }, { name = "python-dateutil", specifier = ">2,<3" }, + { name = "rich", specifier = ">=13" }, { name = "typeguard", specifier = ">=4,<5" }, { name = "types-paramiko", specifier = ">=2.7,<6" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, @@ -1467,7 +1465,6 @@ requires-dist = [ [package.metadata.requires-dev] dev = [ - { name = "click", specifier = ">=8.2" }, { name = "coverage", specifier = ">=7.7.1,<8" }, { name = "freezegun", specifier = ">=1.5.5" }, { name = "ipdb" }, @@ -1494,7 +1491,6 @@ docs = [ { name = "zensical", specifier = ">=0.0.34" }, ] test = [ - { name = "click", specifier = ">=8.2" }, { name = "coverage", specifier = ">=7.7.1,<8" }, { name = "freezegun", specifier = ">=1.5.5" }, { name = "mypy", specifier = "==1.17.1" }, From d1de5cabd0963779e37240f9257a15c4604978da Mon Sep 17 00:00:00 2001 From: "Axel H." Date: Sun, 26 Jul 2026 01:10:43 +0200 Subject: [PATCH 2/2] feat(cli): live hierarchical progress tree and informative error handling Add a live, hierarchical progress tree for deploys under a TTY: each phase (connect, prepare, execute) and each operation renders as a tree of nodes with per-host status, and host-attributed log/echo output nests under the relevant host node. Non-TTY, --json and --debug runs keep the flat log output. Also flips the host print prefix from `[name]` to a styled `name` and adds an `operation_host_skipped` state callback so skipped hosts are represented. Fixes three TTY output-loss regressions reported on review: - `--serial` / `--no-wait` lost all execution output: these paths drive hosts directly and never fire `operation_start` / `operation_end`, so the operation tree node was never created and every host line was dropped. The op node is now created lazily on first host activity, and still-running op nodes are finalised (succeed/fail from their children) when a region rotates or the tree stops. - `--diff` dropped the diff unless `-v`: host-attributed INFO output (the diff header and body) was gated behind verbose mode. The gate now only hides known lifecycle noise (Connected / Ready / Disconnected / noop); real deploy output and diffs always route to the host node. - Host-attributed deploy INFO logs were dropped under a TTY for the same reason and are now shown. Includes unit coverage for the lazy-node creation, pending-op finalisation and the lifecycle-noise gate. --- src/pyinfra/api/connect.py | 4 +- src/pyinfra/api/facts.py | 7 +- src/pyinfra/api/host.py | 38 +- src/pyinfra/api/operations.py | 19 +- src/pyinfra/api/renderable.py | 38 ++ src/pyinfra/api/state.py | 4 + src/pyinfra/api/util.py | 2 +- src/pyinfra/operations/files.py | 17 +- src/pyinfra/operations/util/files.py | 28 +- src/pyinfra/progress.py | 190 ++++--- src/pyinfra_cli/cli.py | 202 +++++--- src/pyinfra_cli/console.py | 77 ++- src/pyinfra_cli/exceptions.py | 12 +- src/pyinfra_cli/log.py | 81 ++- src/pyinfra_cli/main.py | 1 + src/pyinfra_cli/prints.py | 205 ++++---- src/pyinfra_cli/progress.py | 670 +++++++++++++++++++++++++ src/pyinfra_cli/renderables.py | 68 +++ src/pyinfra_cli/routing.py | 139 +++++ src/pyinfra_cli/util.py | 28 +- tests/end-to-end/conftest.py | 2 +- tests/end-to-end/test_e2e_local.py | 32 +- tests/end-to-end/test_e2e_ssh.py | 28 +- tests/test_cli/test_cli_log.py | 70 +++ tests/test_cli/test_cli_progress.py | 221 ++++++++ tests/test_cli/test_cli_renderables.py | 61 +++ tests/test_operations_utils.py | 39 +- 27 files changed, 1916 insertions(+), 367 deletions(-) create mode 100644 src/pyinfra/api/renderable.py create mode 100644 src/pyinfra_cli/progress.py create mode 100644 src/pyinfra_cli/renderables.py create mode 100644 src/pyinfra_cli/routing.py create mode 100644 tests/test_cli/test_cli_log.py create mode 100644 tests/test_cli/test_cli_progress.py create mode 100644 tests/test_cli/test_cli_renderables.py diff --git a/src/pyinfra/api/connect.py b/src/pyinfra/api/connect.py index ca119487e..23113abd9 100644 --- a/src/pyinfra/api/connect.py +++ b/src/pyinfra/api/connect.py @@ -24,7 +24,7 @@ def connect_all(state: "State"): greenlet_to_host = {state.pool.spawn(host.connect): host for host in hosts} - with progress_spinner(greenlet_to_host.values()) as progress: + with progress_spinner(greenlet_to_host.values(), prefix_message="Connecting") as progress: for greenlet in gevent.iwait(greenlet_to_host.keys()): host = greenlet_to_host[greenlet] progress(host) @@ -57,7 +57,7 @@ def disconnect_all(state: "State"): for host in state.activated_hosts # only hosts we connected to please! } - with progress_spinner(greenlet_to_host.values()) as progress: + with progress_spinner(greenlet_to_host.values(), prefix_message="Disconnecting") as progress: for greenlet in gevent.iwait(greenlet_to_host.keys()): host = greenlet_to_host[greenlet] progress(host) diff --git a/src/pyinfra/api/facts.py b/src/pyinfra/api/facts.py index f6e4b74a5..86e3fb569 100644 --- a/src/pyinfra/api/facts.py +++ b/src/pyinfra/api/facts.py @@ -176,7 +176,12 @@ def get_host_fact(host, *args, **kwargs): results = {} - with progress_spinner(greenlet_to_host.values()) as progress: + fact_cls = args[0] if args else None + fact_name = getattr(fact_cls, "name", None) or getattr(fact_cls, "__name__", "facts") + + with progress_spinner( + greenlet_to_host.values(), prefix_message=f"Gathering {fact_name}" + ) as progress: for greenlet in gevent.iwait(greenlet_to_host.keys()): host = greenlet_to_host[greenlet] results[host] = greenlet.get() diff --git a/src/pyinfra/api/host.py b/src/pyinfra/api/host.py index 64a99fd6c..ac5d541ca 100644 --- a/src/pyinfra/api/host.py +++ b/src/pyinfra/api/host.py @@ -16,6 +16,7 @@ from pyinfra import logger from pyinfra.api.output import format_text +from pyinfra.api.renderable import Diff, OutputBlock from pyinfra.connectors.base import BaseConnector from pyinfra.connectors.util import CommandOutput, remove_any_sudo_askpass_file from pyinfra.context import ctx_config @@ -198,19 +199,50 @@ def host_data(self): def group_data(self): return self.inventory.get_groups_data(self.groups) + def _styled_name(self, *args, **kwargs) -> str: + # Dim any "@connector/" prefix so the host name stands out. + name = self.name + if name.startswith("@") and "/" in name: + connector, _, rest = name.partition("/") + return ( + f"{format_text(f'{connector}/', 'bright_black')}" + f"{format_text(rest, *args, **kwargs)}" + ) + return format_text(name, *args, **kwargs) + @property def print_prefix(self) -> str: if self.nested_executing_op_hash: - return f"{format_text('')}[{format_text(self.name, bold=True)}] {format_text('nested', 'blue')}{self.print_prefix_padding} " + return f"{self._styled_name('cyan', bold=True)} {format_text('nested', 'blue')}{self.print_prefix_padding} " - return f"{format_text('')}[{format_text(self.name, bold=True)}]{self.print_prefix_padding} " + return f"{self._styled_name('cyan', bold=True)}{self.print_prefix_padding} " def style_print_prefix(self, *args, **kwargs) -> str: - return f"{format_text('')}[{format_text(self.name, *args, **kwargs)}]{self.print_prefix_padding} " + return f"{self._styled_name(*args, **kwargs)}{self.print_prefix_padding} " def log(self, message: str, log_func: Callable[[str], Any] = logger.info) -> None: log_func(f"{self.print_prefix}{message}") + def log_rich( + self, + descriptor: OutputBlock, + log_func: Callable[..., Any] = logger.info, + ) -> None: + """Emit a rich-free output descriptor as a single log record. + + The descriptor and this host's name are attached to the record via + ``extra``; the CLI turns the descriptor into a rich renderable and routes + it to this host's tree node (or the console in flat mode). Keeping the + descriptor plain lets the core stay decoupled from any rendering library. + """ + log_func("", extra={"pyinfra_rich": descriptor, "pyinfra_host": self.name}) + + def log_diff(self, diff_text: str) -> None: + """Log a whole (plain) diff block for syntax-highlighted rendering.""" + if not diff_text: + return + self.log_rich(Diff(diff_text)) + def log_styled( self, message: str, log_func: Callable[[str], Any] = logger.info, **kwargs ) -> None: diff --git a/src/pyinfra/api/operations.py b/src/pyinfra/api/operations.py index 4645c6c21..947f45096 100644 --- a/src/pyinfra/api/operations.py +++ b/src/pyinfra/api/operations.py @@ -39,6 +39,7 @@ def run_host_op(state: State, host: Host, op_hash: str) -> bool: if op_hash not in state.ops[host]: logger.info(f"{host.print_prefix}{format_text('Skipped', 'blue')}") + state.trigger_callbacks("operation_host_skipped", host, op_hash) return True op_meta = state.get_op_meta(op_hash) @@ -283,7 +284,7 @@ def _run_serial_ops(state: State): for host in list(state.inventory.get_active_hosts()): host_operations = product([host], state.get_op_order()) - with progress_spinner(host_operations) as progress: + with progress_spinner(host_operations, prefix_message=f"Running ({host.name})") as progress: try: _run_host_ops( state, @@ -300,7 +301,7 @@ def _run_no_wait_ops(state: State): """ hosts_operations = product(state.inventory.get_active_hosts(), state.get_op_order()) - with progress_spinner(hosts_operations) as progress: + with progress_spinner(hosts_operations, prefix_message="Running operations") as progress: # Spawn greenlet for each host to run *all* ops if state.pool is None: raise PyinfraError("No pool found on state.") @@ -326,10 +327,14 @@ def _run_single_op(state: State, op_hash: str): op_meta = state.get_op_meta(op_hash) log_operation_start(op_meta) + op_name = ", ".join(op_meta.names) if op_meta.names else "operation" + failed_hosts = set() if op_meta.global_arguments["_serial"]: - with progress_spinner(state.inventory.get_active_hosts()) as progress: + with progress_spinner( + state.inventory.get_active_hosts(), prefix_message=op_name + ) as progress: # For each host, run the op for host in state.inventory.get_active_hosts(): result = _run_host_op_with_context(state, host, op_hash) @@ -349,7 +354,7 @@ def _run_single_op(state: State, op_hash: str): batches = [hosts[i : i + parallel] for i in range(0, len(hosts), parallel)] for batch in batches: - with progress_spinner(batch) as progress: + with progress_spinner(batch, prefix_message=op_name) as progress: # Spawn greenlet for each host if state.pool is None: raise PyinfraError("No pool found on state.") @@ -368,11 +373,13 @@ def _run_single_op(state: State, op_hash: str): if not greenlet.get(): failed_hosts.add(host) + # Signal the operation end first so progress handlers can finalise its + # display before fail_hosts potentially prompts or raises. + state.trigger_callbacks("operation_end", op_hash) + # Now all the batches/hosts are complete, fail any failures state.fail_hosts(failed_hosts) - state.trigger_callbacks("operation_end", op_hash) - def run_ops(state: State, serial: bool = False, no_wait: bool = False): """ diff --git a/src/pyinfra/api/renderable.py b/src/pyinfra/api/renderable.py new file mode 100644 index 000000000..cc965e9db --- /dev/null +++ b/src/pyinfra/api/renderable.py @@ -0,0 +1,38 @@ +""" +Rich-free output descriptors for structured CLI output. + +Operations and facts describe rich output (diffs, syntax-highlighted code +blocks, ...) with these plain dataclasses via :meth:`pyinfra.api.host.Host.log_rich`. +The core stays decoupled from any rendering library — the CLI layer +(``pyinfra_cli``) maps each descriptor to a Rich renderable via its registry. + +New descriptor types only need a subclass here plus a CLI-side renderer +registration; the routing/rendering pipeline is generic. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class OutputBlock: + """Base marker for a structured, rich-free output descriptor.""" + + +@dataclass(frozen=True) +class Diff(OutputBlock): + """A unified-diff block (rendered with the ``diff`` lexer).""" + + text: str + + +@dataclass(frozen=True) +class CodeBlock(OutputBlock): + """A syntax-highlighted code block. + + ``lexer`` is a Pygments lexer name (``sql``, ``yaml``, ``json``, ...). + """ + + text: str + lexer: str = "text" diff --git a/src/pyinfra/api/state.py b/src/pyinfra/api/state.py index 083413c17..66ca71562 100644 --- a/src/pyinfra/api/state.py +++ b/src/pyinfra/api/state.py @@ -70,6 +70,10 @@ def operation_start(state: State, op_hash): def operation_host_start(state: State, host: Host, op_hash): pass + @staticmethod + def operation_host_skipped(state: State, host: Host, op_hash): + pass + @staticmethod def operation_host_success(state: State, host: Host, op_hash, retry_count: int = 0): pass diff --git a/src/pyinfra/api/util.py b/src/pyinfra/api/util.py index 432523cf5..681613b93 100644 --- a/src/pyinfra/api/util.py +++ b/src/pyinfra/api/util.py @@ -201,7 +201,7 @@ def print_host_combined_output(host: Host, output: CommandOutput) -> None: def log_operation_start( - op_meta: StateOperationMeta, op_types: list | None = None, prefix: str = "--> " + op_meta: StateOperationMeta, op_types: list | None = None, prefix: str = "" ) -> None: op_types = op_types or [] if op_meta.global_arguments["_serial"]: diff --git a/src/pyinfra/operations/files.py b/src/pyinfra/operations/files.py index 01d38ef22..85c6d50b4 100644 --- a/src/pyinfra/operations/files.py +++ b/src/pyinfra/operations/files.py @@ -66,7 +66,7 @@ MetadataTimeField, adjust_regex, ensure_mode_int, - generate_color_diff, + generate_diff, get_timestamp, sed_delete, sed_replace, @@ -505,9 +505,7 @@ def line( elif present_lines and not present: if state.config.DIFF: host.log(f"Will Remove lines in {format_text(path, bold=True)}", logger.info) - for line in generate_color_diff(present_lines, []): - logger.info(" %s", line) - logger.info("") + host.log_diff("\n".join(generate_diff(present_lines, []))) yield sed_delete( path, match_line, @@ -525,8 +523,7 @@ def line( if state.config.DIFF: host.log(f"Will replace lines in {format_text(path, bold=True)}", logger.info) new_lines = [re.sub(match_line, replace, line) for line in present_lines] - for line in generate_color_diff(present_lines, new_lines): - logger.info(" %s", line) + host.log_diff("\n".join(generate_diff(present_lines, new_lines))) yield sed_replace_command else: host.noop(f'line "{replace or line}" exists in {path}') @@ -1159,11 +1156,9 @@ def put( with get_file_io(src, "r") as f: desired_lines = f.readlines() - for line in generate_color_diff([], desired_lines): - logger.info(" %s", line) + host.log_diff("\n".join(generate_diff([], desired_lines))) except UnicodeDecodeError: logger.info("Binary file uploaded") - logger.info("") yield FileUploadCommand( local_file, @@ -1208,9 +1203,7 @@ def put( with get_file_io(src, "r") as f: desired_lines = f.readlines() - for line in generate_color_diff(current_lines, desired_lines): - logger.info(" %s", line) - logger.info("") + host.log_diff("\n".join(generate_diff(current_lines, desired_lines))) yield FileUploadCommand( local_file, diff --git a/src/pyinfra/operations/util/files.py b/src/pyinfra/operations/util/files.py index 3b2bcdfd3..0d5901d7c 100644 --- a/src/pyinfra/operations/util/files.py +++ b/src/pyinfra/operations/util/files.py @@ -261,9 +261,14 @@ def strip_regex_anchors(line: str) -> str: return line -def generate_color_diff( - current_lines: list[str], desired_lines: list[str] -) -> Generator[str, None, None]: +def generate_diff(current_lines: list[str], desired_lines: list[str]) -> Generator[str, None, None]: + """Yield a plain (uncoloured) unified diff between two sets of lines. + + Removed lines are prefixed ``"- "``, added lines ``"+ "``, context lines + ``" "`` and hunks with ``"@@ ... @@"`` headers. Kept plain so callers can + render it however they like (e.g. syntax highlighting). + """ + def _format_range_unified(start: int, stop: int) -> str: beginning = start + 1 # lines start numbering with one length = stop - start @@ -286,7 +291,20 @@ def _format_range_unified(start: int, stop: int) -> str: continue if tag in {"replace", "delete"}: for line in current_lines[i1:i2]: - yield format_text("- " + line.rstrip(), "red") + yield "- " + line.rstrip() if tag in {"replace", "insert"}: for line in desired_lines[j1:j2]: - yield format_text("+ " + line.rstrip(), "green") + yield "+ " + line.rstrip() + + +def generate_color_diff( + current_lines: list[str], desired_lines: list[str] +) -> Generator[str, None, None]: + """Yield a unified diff with ``format_text`` colouring (- red, + green).""" + for line in generate_diff(current_lines, desired_lines): + if line.startswith("- "): + yield format_text(line, "red") + elif line.startswith("+ "): + yield format_text(line, "green") + else: + yield line diff --git a/src/pyinfra/progress.py b/src/pyinfra/progress.py index 1b5b37f17..277b3db3a 100644 --- a/src/pyinfra/progress.py +++ b/src/pyinfra/progress.py @@ -1,130 +1,116 @@ -import math +from __future__ import annotations + import os -import platform -import sys -from collections import deque from contextlib import contextmanager -import gevent -from gevent.event import Event - -from pyinfra.api.output import is_output_active - -IS_WINDOWS = platform.system() == "Windows" - -WAIT_TIME = 1 / 5 -WAIT_CHARS = deque(("-", "/", "|", "\\")) - -# Hacky way of getting terminal size (so can clear lines) -# Source: http://stackoverflow.com/questions/566746 -IS_TTY = sys.stdout.isatty() and sys.stderr.isatty() -TERMINAL_WIDTH = 0 - -if IS_TTY: - try: - TERMINAL_WIDTH = os.get_terminal_size().columns - except AttributeError: - if not IS_WINDOWS: - terminal_size = os.popen("stty size", "r").read().split() - if len(terminal_size) == 2: - TERMINAL_WIDTH = int(terminal_size[1]) +from typing import TYPE_CHECKING, Any +from rich.errors import LiveError +from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn -def _print_spinner(stop_event, progress_queue): - if not IS_TTY or os.environ.get("PYINFRA_PROGRESS") == "off": - return +from pyinfra.api.output import get_console, is_output_active - progress = "" - text = "" +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Iterator - while True: - # Stop when asked too - if stop_event.is_set(): - break +# A single shared Progress instance is reused for the whole run so that +# concurrent/nested phases (connect, prepare, execute, ...) each get their own +# bar within one live display. Per-host log lines printed via the shared +# console appear *above* the live bars automatically. +# +# The module-level refcount is mutated from multiple greenlets without a lock; +# this is safe because greenlets are cooperative and ``auto_refresh=False`` +# means there is no background refresh thread racing the mutations. +_progress: Progress | None = None +_active_spinners = 0 - WAIT_CHARS.rotate(1) - try: - progress = progress_queue[-1] - except IndexError: - pass +def _get_progress() -> Progress: + global _progress + if _progress is None: + _progress = Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TaskProgressColumn(), + TextColumn("{task.completed}/{task.total}"), + console=get_console(), + transient=True, + auto_refresh=False, + ) + return _progress - text = f" {' '.join((WAIT_CHARS[0], progress))}" - text = f"{text}\r" - sys.stderr.write(text) - sys.stderr.flush() +def _spinner_enabled() -> bool: + # Only render when in CLI mode and not explicitly disabled. + return is_output_active() and os.environ.get("PYINFRA_PROGRESS") != "off" - # In pyinfra_cli's __main__ we set stdout & stderr to be line buffered, - # so write this escape code (clear line) into the buffer but don't flush, - # such that any next print/log/etc clear the line first. - if not IS_WINDOWS: - sys.stderr.write("\033[K") - stop_event.wait(timeout=WAIT_TIME) +def _noop_progress(complete_item: Any) -> None: + pass @contextmanager -def progress_spinner(items, prefix_message=None): - # If there's no current state context we're not in CLI mode, so just return a noop - # handler and exit. - if not is_output_active(): - yield lambda complete_item: None +def progress_spinner( + items: Iterable[Any], + prefix_message: str | None = None, +) -> Iterator[Callable[[Any], None]]: + """ + Display a Rich progress bar while ``items`` are completed. + + Yields a ``progress(complete_item)`` callback; callers may ignore it (using + the bar purely as a "busy" indicator). Multiple/nested calls share a single + live display, each contributing its own bar. The display is refreshed + manually (``auto_refresh=False``) from the callback to stay well-behaved + under gevent (no background refresh greenlet). + """ + if not _spinner_enabled(): + yield _noop_progress return + global _active_spinners, _progress + if not isinstance(items, set): items = set(items) total_items = len(items) - stop_event = Event() - - def make_progress_message(include_items=True): - message_bits = [] - - # If we only have 1 item, don't show % - if total_items > 1: - percentage_complete = 0 - complete = total_items - len(items) - percentage_complete = int(math.floor(complete / total_items * 100)) - message_bits.append( - f"{percentage_complete}% ({complete}/{total_items})", - ) - - if prefix_message: - message_bits.append(prefix_message) + progress_bar = _get_progress() - if include_items and items: - # Plus 3 for the " - " joining below - message_length = sum((len(message) + 3) for message in message_bits) - # -8 for padding left+right, -2 for {} wrapping - items_allowed_width = TERMINAL_WIDTH - 10 - message_length - - if items_allowed_width > 0: - items_string = f"{{{', '.join(f'{i}' for i in items)}}}" - if len(items_string) >= items_allowed_width: - # -3 for the ... - items_string = f"{items_string[: items_allowed_width - 3]}...}}" - - message_bits.append(items_string) - - return " - ".join(message_bits) - - progress_queue = deque((make_progress_message(),)) - - def progress(complete_item): + if _active_spinners == 0: + try: + progress_bar.start() + except LiveError: + # Another live display owns the shared console (e.g. the CLI's + # live progress tree) — rich only allows one at a time. + _progress = None + yield _noop_progress + return + _active_spinners += 1 + + description = prefix_message or "Working" + task_id = progress_bar.add_task(description, total=total_items) + progress_bar.refresh() + + def progress(complete_item: Any) -> None: if complete_item not in items: raise ValueError( f"Invalid complete item: {complete_item} not in {items}", ) - items.remove(complete_item) - progress_queue.append(make_progress_message()) - - # Kick off the spinner greenlet - spinner_greenlet = gevent.spawn(_print_spinner, stop_event, progress_queue) - - # Yield allowing the actual code the spinner waits for to run - yield progress + progress_bar.update(task_id, advance=1) + progress_bar.refresh() - # Finally, stop the spinner - stop_event.set() - spinner_greenlet.join() + try: + yield progress + finally: + # Decrement first so the display is always stopped even if the task + # removal fails. + _active_spinners -= 1 + try: + progress_bar.remove_task(task_id) + progress_bar.refresh() + finally: + if _active_spinners == 0: + progress_bar.stop() + # Drop the instance so a fresh one is created for the next run + # (important for long-lived processes / tests). + _progress = None diff --git a/src/pyinfra_cli/cli.py b/src/pyinfra_cli/cli.py index 3194b6c63..b5d735678 100644 --- a/src/pyinfra_cli/cli.py +++ b/src/pyinfra_cli/cli.py @@ -26,7 +26,9 @@ from pyinfra.api.output import format_text from .commands import get_facts_and_args, get_func_and_args +from . import routing from .console import console, stdout_console +from .progress import DeployProgress, is_tree_active, step from .exceptions import CliError, UnexpectedExternalError, UnexpectedInternalError, WrappedError from .inventory import make_inventory from .log import setup_logging @@ -380,8 +382,11 @@ def cli( # Re-raise any unexpected internal exceptions as UnexpectedInternalError raise UnexpectedInternalError(e) finally: + # Stop routing host output into the (stopped) live tree: disconnect + # notices (e.g. docker image IDs) must reach the console. + routing.set_tree(None) if ctx_state.isset() and state.initialised: - logger.info("--> Disconnecting from hosts...") + logger.info("Disconnecting from hosts...") # Triggers any executor disconnect requirements disconnect_all(state) @@ -426,12 +431,6 @@ def _main( debug_operations: bool, json_output: bool = False, ): - # In JSON mode keep the spinner quiet so stdout stays pure JSON. Do not - # force --yes: a JSON run must be able to diff a host without mutating - # it. Applying still requires an explicit --yes; without it the proposed - # changes are emitted as JSON instead of blocking on a confirm prompt. - if json_output: - environ.setdefault("PYINFRA_PROGRESS", "off") # Setup working directory # if chdir: @@ -452,27 +451,45 @@ def _main( config = Config() ctx_config.set(config) + # Decide whether to use the hierarchical live tree (TTY, not JSON, not + # --debug). It renders at every verbosity level: host log/echo lines are + # routed into the host's tree node (see pyinfra_cli.routing) so nothing + # interleaves with the live region; verbosity only controls how much + # detail the core emits. In other modes (piped, JSON, --debug) we fall + # back to plain flat logs. The low-level progress bars are disabled unless + # the user explicitly exports PYINFRA_PROGRESS: the tree replaces them, in + # flat modes they'd fight the log stream, and in JSON mode stdout must + # stay pure JSON. + environ.setdefault("PYINFRA_PROGRESS", "off") + routing.reset_host_errors() + tree = None + if is_tree_active(json_output) and not (debug or debug_all): + tree = DeployProgress(state, verbose=verbosity > 0) + # NOTE: registered as a state callback after state.init() below. + routing.set_tree(tree) + # Update Config & Override Data # - config = _set_config( - config, - config_filename, - sudo, - sudo_user, - use_sudo_password, - use_sudo_login, - same_sudo_password, - su_user, - dzdo, - dzdo_user, - parallel, - shell_executable, - fail_percent, - yes, - diff, - retry, - retry_delay, - ) + with step(tree, "Loading config"): + config = _set_config( + config, + config_filename, + sudo, + sudo_user, + use_sudo_password, + use_sudo_login, + same_sudo_password, + su_user, + dzdo, + dzdo_user, + parallel, + shell_executable, + fail_percent, + yes, + diff, + retry, + retry_delay, + ) if ssh_password_prompt: ssh_password = getpass("SSH password: ") @@ -492,21 +509,29 @@ def _main( # Load up the inventory from the filesystem # - logger.info("--> Loading inventory...") - inventory = make_inventory( - inventory, - cwd=state.cwd, - override_data=override_data, - group_data_directories=group_data, - ) - ctx_inventory.set(inventory) + with step(tree, "Loading inventory"): + logger.info("Loading inventory...") + inventory = make_inventory( + inventory, + cwd=state.cwd, + override_data=override_data, + group_data_directories=group_data, + ) + ctx_inventory.set(inventory) + + # Now that we have inventory, apply --limit/--exclude config override + initial_limit = _apply_inventory_limit(inventory, limit) + initial_limit = _apply_inventory_exclude(inventory, initial_limit, exclude) - # Now that we have inventory, apply --limit/--exclude config override - initial_limit = _apply_inventory_limit(inventory, limit) - initial_limit = _apply_inventory_exclude(inventory, initial_limit, exclude) + # Initialise the state + state.init(inventory, config, initial_limit=initial_limit) - # Initialise the state - state.init(inventory, config, initial_limit=initial_limit) + # Register the inventory host names for log/echo host attribution. + routing.set_host_names(host.name for host in inventory) + + # Now that state is initialised, register the live-tree callback handler. + if tree is not None: + state.add_callback_handler(tree) if command == CliCommands.DEBUG_INVENTORY: if json_output: @@ -517,22 +542,45 @@ def _main( # Connect to the hosts & start handling the user commands # - logger.info("--> Connecting to hosts...") + logger.info("Connecting to hosts...") state.set_stage(StateStage.Connect) - connect_all(state) + if tree is not None: + with tree: + connect_all(state) + else: + connect_all(state) state.set_stage(StateStage.Prepare) - can_diff, state, config = _handle_commands( - state, config, command, original_operations, operations, json_output=json_output - ) + try: + if tree is not None: + with tree: + can_diff, state, config = _handle_commands( + state, config, command, original_operations, operations, json_output=json_output + ) + else: + can_diff, state, config = _handle_commands( + state, config, command, original_operations, operations, json_output=json_output + ) + except PyinfraError: + # e.g. "No hosts remaining!" when every host failed during prepare: + # show what failed before the error propagates. + if state.failed_hosts and not json_output: + _print_failed_hosts(state) + raise + + # Failure prompts are deferred during Prepare (hosts evaluate the deploy + # in parallel; prompting mid-phase interleaves with other hosts' output). + # Now the phase is complete, show what failed and ask once. + if state.failed_hosts and yes is False and not json_output: + if not _confirm_failed_hosts(state, "One or more hosts failed, continue?"): + _exit() # Print proposed changes, execute unless --dry, and exit # if can_diff and not json_output: if yes: - logger.info("--> Skipping change detection") + logger.info("Skipping change detection") else: - logger.info("--> Detected changes:") print_meta(state) console.print( """ @@ -566,11 +614,17 @@ def _main( if not _do_confirm("Detected changes displayed above, skip this step with -y"): _exit() - logger.info("--> Beginning operation run...") + logger.info("Beginning operation run...") state.set_stage(StateStage.Execute) - run_ops(state, serial=serial, no_wait=no_wait) + if tree is not None: + with tree: + run_ops(state, serial=serial, no_wait=no_wait) + # The live display is over; host output from now on (results, + # disconnect notices) streams straight to the console. + routing.set_tree(None) + else: + run_ops(state, serial=serial, no_wait=no_wait) - logger.info("--> Results:") state.set_stage(StateStage.Disconnect) if json_output: print_run_json(state, dry=False) @@ -591,6 +645,35 @@ def _do_confirm(msg: str) -> bool: return Confirm.ask(" Execute?", console=console, default=True) +def _print_failed_hosts(state: State) -> None: + """List the failed hosts with their first recorded error message.""" + console.print() + console.print("Failed hosts:", style="bold red") + host_errors = routing.get_host_errors() + for host in sorted(state.failed_hosts, key=lambda h: h.name): + errors = host_errors.get(host.name) or [] + label = routing.host_label(host.name, base_style="red", prefix=" ✗ ") + if errors: + label.append(f" — {errors[0]}", style="red") + console.print(label) + + +def _confirm_failed_hosts(state: State, msg: str) -> bool: + """Show which hosts failed (and why) then ask whether to continue. + + Pauses the live progress tree (if running) around the prompt so it doesn't + fight the interactive input, and resumes it afterwards. + """ + tree = routing.get_tree() + paused = tree.pause() if tree is not None else False + try: + _print_failed_hosts(state) + return _do_confirm(msg) + finally: + if paused and tree is not None: + tree.resume() + + # Setup # def _setup_log_level(debug, debug_all): @@ -719,7 +802,7 @@ def _set_config( retry, retry_delay, ): - logger.info("--> Loading config...") + logger.info("Loading config...") # Load up any config.py from the filesystem if state.cwd: @@ -811,9 +894,16 @@ def _set_fail_prompts(state: State, config: Config) -> None: config.FAIL_PERCENT = 0 def should_raise_failed_hosts(state: State) -> bool: + if state.current_stage == StateStage.Prepare: + # Hosts prepare in parallel: prompting now would interleave with + # the other hosts' output. Continue silently; one aggregated + # prompt is shown after the phase completes (see _main). + return False if state.current_stage == StateStage.Connect: - return not _do_confirm("One of more hosts failed to connect, continue?") - return not _do_confirm("One of more hosts failed, continue?") + return not _confirm_failed_hosts( + state, "One or more hosts failed to connect, continue?" + ) + return not _confirm_failed_hosts(state, "One or more hosts failed, continue?") state.should_raise_failed_hosts = should_raise_failed_hosts @@ -867,7 +957,7 @@ def _apply_inventory_exclude( # def _handle_commands(state, config, command, original_operations, operations, json_output=False): if command is CliCommands.FACT: - logger.info("--> Gathering facts...") + logger.info("Gathering facts...") state, fact_data = _run_fact_operations(state, config, operations) if json_output: print_facts_json(fact_data) @@ -878,16 +968,16 @@ def _handle_commands(state, config, command, original_operations, operations, js can_diff = True if command == CliCommands.SHELL: - logger.info("--> Preparing exec operation...") + logger.info("Preparing exec operation...") state = _prepare_exec_operations(state, config, operations) can_diff = False elif command == CliCommands.DEPLOY_FILES: - logger.info("--> Preparing operation files...") + logger.info("Preparing operation files...") state, config, operations = _prepare_deploy_operations(state, config, operations) elif command == CliCommands.FUNC: - logger.info("--> Preparing operation func...") + logger.info("Preparing operation func...") state, kwargs = _prepare_func_operations( state, config, diff --git a/src/pyinfra_cli/console.py b/src/pyinfra_cli/console.py index 8fb719eae..b7337bdb6 100644 --- a/src/pyinfra_cli/console.py +++ b/src/pyinfra_cli/console.py @@ -1,5 +1,6 @@ """ -Shared Rich consoles and Click-compatible output adapters for the CLI. +Shared Rich consoles and output adapters for the CLI (keeping the legacy +``click.style``/``click.echo`` call signatures used across the core library). pyinfra keeps all human-facing output on **stderr** and reserves **stdout** for machine-readable (``--json``) payloads. The core library styles/echoes text @@ -8,22 +9,69 @@ from __future__ import annotations +from functools import cache from typing import Any -from rich.console import Console +from rich.console import Console, RenderableType +from rich.padding import Padding +from rich.syntax import Syntax from rich.text import Text -from pyinfra.api.output import get_console, set_console +from pyinfra.api.output import get_console + +from . import routing # Human-facing console (logs, tables, prompts, spinner) → stderr. # Reuse the core shared console so the progress spinner and logging write to the # same Console instance (avoids Live-region corruption). console = get_console() -set_console(console) # Machine-readable console (``--json`` payloads) → stdout, no styling. stdout_console = Console(highlight=False, soft_wrap=True, markup=False, emoji=False) +# Detached console used ONLY to render styled text to ANSI in format_text(). +# It must not be the shared console: capturing on a console with an active +# Live region would embed the whole re-rendered live frame in the capture. +_capture_console = Console( + stderr=True, + highlight=False, + soft_wrap=True, + markup=False, + emoji=False, +) + + +def diff_renderable(diff_text: str, indent: int = 0) -> RenderableType: + """Render a plain unified diff as a syntax-highlighted block. + + ``ansi_dark`` uses the terminal's ANSI palette (theme-friendly) and + ``background_color="default"`` avoids a solid block against the terminal + background — important inside the live region. + """ + syntax = Syntax( + diff_text, + "diff", + background_color="default", + theme="ansi_dark", + word_wrap=True, + ) + if indent: + return Padding(syntax, (0, 0, 0, indent)) + return syntax + + +@cache +def _style_codes(style: str) -> tuple[str, str]: + """ANSI (prefix, suffix) escape codes for a Rich style string, cached. + + ``format_text`` runs in hot paths (per command-output/log/diff line), so + the style→ANSI rendering is done once per style instead of per call. + """ + with _capture_console.capture() as capture: + _capture_console.print(Text("|", style=style), end="") + prefix, _, suffix = capture.get().partition("|") + return prefix, suffix + def format_text(text: str, fg: str | None = None, *, bold: bool = False, **kwargs: Any) -> str: """ @@ -32,7 +80,8 @@ def format_text(text: str, fg: str | None = None, *, bold: bool = False, **kwarg Mirrors the legacy ``click.style`` signature (positional foreground color + ``bold=``) used across the core library, but renders via Rich so styling is consistent with the rest of the CLI output. Colour names (``red``, - ``green``, ...) are passed straight through to Rich. + ``green``, ...) are passed straight through to Rich; other styling kwargs + are ignored. """ style_bits = [] if fg is not None: @@ -43,10 +92,8 @@ def format_text(text: str, fg: str | None = None, *, bold: bool = False, **kwarg if not style_bits: return text - rich_text = Text(text, style=" ".join(style_bits)) - with console.capture() as capture: - console.print(rich_text, end="") - return capture.get() + prefix, suffix = _style_codes(" ".join(style_bits)) + return f"{prefix}{text}{suffix}" def echo(message: Any = None, *, err: bool = False, nl: bool = True, **kwargs: Any) -> None: @@ -56,7 +103,19 @@ def echo(message: Any = None, *, err: bool = False, nl: bool = True, **kwargs: A ``err=True`` targets the (default) human stderr console; ``err=False`` targets stdout. ``nl=False`` suppresses the trailing newline. Text may contain ANSI escape codes already produced by :func:`format_text`. + + When the live progress tree is active, host-attributed messages (command + input/output, transfer notices, ...) are routed into the host's tree node + instead of streaming to the console. """ + if err and isinstance(message, str): + tree = routing.get_tree() + if tree is not None and tree.is_active: + host_name, text = routing.attribute_host(message) + if host_name is not None: + tree.add_host_detail(host_name, text) + return + target = console if err else stdout_console end = "\n" if nl else "" diff --git a/src/pyinfra_cli/exceptions.py b/src/pyinfra_cli/exceptions.py index c8643554b..2fe938c06 100644 --- a/src/pyinfra_cli/exceptions.py +++ b/src/pyinfra_cli/exceptions.py @@ -100,7 +100,7 @@ def show(self) -> None: name = f"{name} in {info.filename} line {info.lineno}" logger.warning( - f"--> {format_text(name, 'red', bold=True)}: {self}", + f"{format_text(name, 'red', bold=True)}: {self}", ) @@ -108,7 +108,7 @@ class CliError(CliException): @override def show(self) -> None: logger.warning( - f"--> {format_text('pyinfra error', 'red', bold=True)}: {self}", + f"{format_text('pyinfra error', 'red', bold=True)}: {self}", ) @@ -123,7 +123,7 @@ def __init__(self, e, filename): @override def show(self) -> None: logger.warning( - "--> {}:\n".format( + "{}:\n".format( format_text( f"An exception occurred in: {self.filename}", "red", @@ -145,7 +145,7 @@ def __init__(self, e): @override def show(self) -> None: console.print( - "--> {}:\n".format( + "{}:\n".format( format_text( "An internal exception occurred", "red", @@ -165,9 +165,9 @@ def show(self) -> None: logger.debug(str(self.exception)) console.print( - f"--> The full traceback has been written to {format_text('pyinfra-debug.log', bold=True)}", + f"The full traceback has been written to {format_text('pyinfra-debug.log', bold=True)}", ) console.print( - "--> If this is unexpected please consider submitting a bug report " + "If this is unexpected please consider submitting a bug report " "on GitHub, for more information run `pyinfra --support`." ) diff --git a/src/pyinfra_cli/log.py b/src/pyinfra_cli/log.py index a64a48e41..74a9768d0 100644 --- a/src/pyinfra_cli/log.py +++ b/src/pyinfra_cli/log.py @@ -6,16 +6,77 @@ from pyinfra import logger, state from pyinfra.context import ctx_state +from . import routing from .console import console, format_text +from .renderables import to_renderable + +# Host-attributed INFO lines that merely restate a phase the tree node already +# conveys (connect/prepare/disconnect lifecycle). Only these are hidden at +# default verbosity; real deploy output (command results, diffs, "Will modify", +# ...) is always routed to the host node so it is never silently dropped. +_LIFECYCLE_INFO_PREFIXES = ("Connected", "Ready:", "Disconnected", "noop:") + + +def _is_lifecycle_noise(text: str) -> bool: + return text.startswith(_LIFECYCLE_INFO_PREFIXES) class LogHandler(logging.Handler): @override def emit(self, record): try: - message = self.format(record) - # ``message`` may already contain ANSI escape codes (from format_text). - console.print(Text.from_ansi(message)) + # Structured "rich" records (from host.log_rich) carry a rich-free + # descriptor + explicit host; render them as a Rich block, bypassing + # the line-oriented formatter. + descriptor = getattr(record, "pyinfra_rich", None) + if descriptor is not None: + host_name = getattr(record, "pyinfra_host", None) + tree = routing.get_tree() + if tree is not None and tree.is_active and host_name is not None: + tree.add_host_renderable(host_name, descriptor) + else: + console.print(to_renderable(descriptor, indent=4)) + return + + # Count warnings here (not in the formatter) so the counter also + # works when messages are routed into the live tree. + if ctx_state.isset() and record.levelno == logging.WARNING: + state.increment_warning_counter() + + message = record.getMessage() + host_name, text = routing.attribute_host(message) + + # Record per-host warnings/errors so failure prompts can show + # which hosts failed and why, in every output mode. + if host_name is not None and record.levelno >= logging.WARNING: + routing.record_host_error(host_name, text) + + tree = routing.get_tree() + if tree is not None: + if host_name is not None and tree.is_active: + # Host warnings/errors always nest under the host's tree + # node. INFO lines are routed too, EXCEPT known lifecycle + # noise (Connected/Ready/...) which the node status already + # conveys — those only show in verbose mode. This ensures + # deploy output and --diff bodies are never dropped. + if ( + record.levelno >= logging.WARNING + or tree.verbose + or not _is_lifecycle_noise(text) + ): + tree.add_host_detail( + host_name, text, is_error=record.levelno >= logging.ERROR + ) + return + if host_name is None and record.levelno < logging.WARNING: + # Non-host INFO lines (phase headers) are dropped: the tree + # already conveys the phases. + return + # Everything else streams to the console: non-host warnings/ + # errors (above the live region) and host lines emitted while + # no live region is running (e.g. disconnect notices). + + console.print(Text.from_ansi(self.format(record))) except Exception: self.handleError(record) @@ -48,10 +109,14 @@ def format(self, record): # We only handle strings here if isinstance(message, str): - if ctx_state.isset() and record.levelno is logging.WARNING: - state.increment_warning_counter() - - if "-->" in message: + # Header lines are top-level phase messages; per-host lines start + # with the host's print prefix and are indented beneath their + # header. Match on the ANSI-stripped prefix (host names may be + # styled). + prefix_host, _ = routing.split_host_prefix(routing.strip_ansi(message)) + is_header = prefix_host is None + + if is_header: if not self.previous_was_header: console.print() else: @@ -60,7 +125,7 @@ def format(self, record): if record.levelno in self.level_to_format: message = self.level_to_format[record.levelno](message) - self.previous_was_header = "-->" in message + self.previous_was_header = is_header return message # If not a string, pass to standard Formatter diff --git a/src/pyinfra_cli/main.py b/src/pyinfra_cli/main.py index 5c3ee3e04..707ee0049 100644 --- a/src/pyinfra_cli/main.py +++ b/src/pyinfra_cli/main.py @@ -6,6 +6,7 @@ import pyinfra from pyinfra.api.output import set_echo, set_formatter +from . import renderables # noqa: F401 (registers built-in descriptor renderers) from .cli import app from .console import console, echo, format_text from .exceptions import CliException diff --git a/src/pyinfra_cli/prints.py b/src/pyinfra_cli/prints.py index f85cfa4c9..a27f83666 100644 --- a/src/pyinfra_cli/prints.py +++ b/src/pyinfra_cli/prints.py @@ -2,21 +2,22 @@ import json import platform -import re import sys from typing import TYPE_CHECKING, Any -from collections.abc import Callable, Iterator +from collections.abc import Iterator from rich.console import Group from rich.json import JSON from rich.padding import Padding from rich.table import Table from rich.text import Text +from rich.tree import Tree -from pyinfra import __version__, logger +from pyinfra import __version__ from pyinfra.api.host import Host from pyinfra.api.output import format_text +from . import routing from .console import console, stdout_console from .util import json_encode @@ -24,13 +25,6 @@ from pyinfra.api.state import State -ANSI_RE = re.compile(r"\033\[((?:\d|;)*)([a-zA-Z])") - - -def _strip_ansi(value): - return ANSI_RE.sub("", value) - - def _get_group_combinations(inventory: Iterator[Host]): group_combinations: dict[tuple, list[Host]] = {} @@ -241,14 +235,14 @@ def print_state_operations(state: State): state_ops = {host: ops for host, ops in state.ops.items() if state.is_host_in_limit(host)} console.print() - console.print("--> Operations:") + console.print("Operations:") console.print(jsonify(state_ops, indent=4, default=json_encode)) console.print() - console.print("--> Operation meta:") + console.print("Operation meta:") console.print(jsonify(state.op_meta, indent=4, default=json_encode)) console.print() - console.print("--> Operation order:") + console.print("Operation order:") console.print() for op_hash in state.get_op_order(): meta = state.op_meta[op_hash] @@ -365,7 +359,7 @@ def print_facts(facts): for name, data in facts.items(): console.print() console.print( - f"--> Fact data for: {format_text(name, bold=True)}", + f"Fact data for: {format_text(name, bold=True)}", ) print_fact(data) @@ -409,54 +403,6 @@ def print_support_info() -> None: ) -def print_rows(rows): - # Go through the rows and work out all the widths in each column - row_column_widths: list[list[int]] = [] - - for _, columns in rows: - if isinstance(columns, str): - continue - - for i, column in enumerate(columns): - if i >= len(row_column_widths): - row_column_widths.append([]) - - # Length of the column (with ansi codes removed) - width = len(_strip_ansi(column.strip())) - row_column_widths[i].append(width) - - # Get the max width of each column and add 4 padding spaces - column_widths = [max(widths) + 4 for widths in row_column_widths] - - # Now print each column, keeping text justified to the widths above - for func, columns in rows: - line = columns - - if not isinstance(columns, str): - justified = [] - - for i, column in enumerate(columns): - stripped = _strip_ansi(column) - desired_width = column_widths[i] - padding = desired_width - len(stripped) - - justified.append( - f"{column}{' '.join('' for _ in range(padding))}", - ) - - line = "".join(justified) - - func(line) - - -def truncate(text, max_length): - if len(text) <= max_length: - return text - - text = text[: max_length - 3] - return f"{text}..." - - def pretty_op_name(op_meta): name = list(op_meta.names)[0] @@ -466,10 +412,17 @@ def pretty_op_name(op_meta): return name +def _split_op_name(name: str) -> tuple[str | None, str]: + """Split a "file.py | Operation" name into (file, operation).""" + if " | " in name: + filename, op_name = name.split(" | ", 1) + return filename, op_name + return None, name + + def print_meta(state: State): - rows: list[tuple[Callable, list[str] | str]] = [ - (logger.info, ["Operation", "Change", "Conditional Change"]), - ] + tree = Tree(Text("Proposed changes", style="bold"), guide_style="dim") + file_branches: dict[str, Any] = {} for op_hash in state.get_op_order(): hosts_in_op = [] @@ -483,32 +436,51 @@ def print_meta(state: State): else: hosts_in_op.append(host.name) - rows.append( - ( - logger.info, - [ - pretty_op_name(state.op_meta[op_hash]), - ( - "-" - if len(hosts_in_op) == 0 - else f"{len(hosts_in_op)} ({truncate(', '.join(sorted(hosts_in_op)), 48)})" - ), - ( - "-" - if len(hosts_maybe_in_op) == 0 - else f"{len(hosts_maybe_in_op)} ({truncate(', '.join(sorted(hosts_maybe_in_op)), 48)})" - ), - ], - ) - ) - - print_rows(rows) + filename, op_name = _split_op_name(pretty_op_name(state.op_meta[op_hash])) + + parent = tree + if filename is not None: + branch = file_branches.get(filename) + if branch is None: + branch = tree.add(Text(filename, style="bold blue")) + file_branches[filename] = branch + parent = branch + + n_change = len(hosts_in_op) + n_maybe = len(hosts_maybe_in_op) + summary = Text(op_name, style="cyan") + if n_change: + summary.append(f" [{n_change} change]", style="green") + if n_maybe: + summary.append(f" [{n_maybe} conditional]", style="yellow") + if not n_change and not n_maybe: + summary.append(" [no change]", style="dim") + + op_branch = parent.add(summary) + for host_name in sorted(hosts_in_op): + op_branch.add(routing.host_label(host_name, base_style="green")) + for host_name in sorted(hosts_maybe_in_op): + label = routing.host_label(host_name, base_style="yellow") + label.append(" (conditional)", style="yellow") + op_branch.add(label) + + console.print(tree) + + +def _result_summary(n_success: int, n_error: int, n_no_change: int) -> Text: + parts = Text() + if n_success: + parts.append(f" {n_success} ✓", style="green") + if n_error: + parts.append(f" {n_error} ✗", style="red") + if n_no_change: + parts.append(f" {n_no_change} –", style="blue") + return parts def print_results(state: State): - rows: list[tuple[Callable, list[str] | str]] = [ - (logger.info, ["Operation", "Hosts", "Success", "Error", "No Change"]), - ] + tree = Tree(Text("Results", style="bold"), guide_style="dim") + file_branches: dict[str, Any] = {} totals = {"hosts": 0, "success": 0, "error": 0, "no_change": 0} @@ -532,37 +504,32 @@ def print_results(state: State): else: hosts_in_op_error.append(host.name) - row = [ - pretty_op_name(state.op_meta[op_hash]), - str(hosts_in_op), - ] - totals["hosts"] += hosts_in_op + totals["success"] += len(hosts_in_op_success) + totals["error"] += len(hosts_in_op_error) + totals["no_change"] += len(hosts_in_op_no_change) + + filename, op_name = _split_op_name(pretty_op_name(state.op_meta[op_hash])) + parent = tree + if filename is not None: + branch = file_branches.get(filename) + if branch is None: + branch = tree.add(Text(filename, style="bold blue")) + file_branches[filename] = branch + parent = branch + + label = Text(op_name, style="red" if hosts_in_op_error else "cyan") + label.append_text( + _result_summary( + len(hosts_in_op_success), len(hosts_in_op_error), len(hosts_in_op_no_change) + ) + ) + op_branch = parent.add(label) + for host_name in sorted(hosts_in_op_error): + op_branch.add(routing.host_label(host_name, base_style="red", prefix="✗ ")) - if hosts_in_op_success: - num_hosts_in_op_success = len(hosts_in_op_success) - row.append(str(num_hosts_in_op_success)) - totals["success"] += num_hosts_in_op_success - else: - row.append("-") - - if hosts_in_op_error: - num_hosts_in_op_error = len(hosts_in_op_error) - row.append(str(num_hosts_in_op_error)) - totals["error"] += num_hosts_in_op_error - else: - row.append("-") - - if hosts_in_op_no_change: - num_hosts_in_op_no_change = len(hosts_in_op_no_change) - row.append(str(num_hosts_in_op_no_change)) - totals["no_change"] += num_hosts_in_op_no_change - else: - row.append("-") - - rows.append((logger.info, row)) - - totals_row = ["Grand total"] + [str(i) if i else "-" for i in totals.values()] - rows.append((logger.info, totals_row)) + grand = Text("Grand total", style="bold") + grand.append_text(_result_summary(totals["success"], totals["error"], totals["no_change"])) + tree.add(grand) - print_rows(rows) + console.print(tree) diff --git a/src/pyinfra_cli/progress.py b/src/pyinfra_cli/progress.py new file mode 100644 index 000000000..48506f4a8 --- /dev/null +++ b/src/pyinfra_cli/progress.py @@ -0,0 +1,670 @@ +""" +Hierarchical live progress renderer for deploys. + +Renders a tree of phases (setup steps, Connecting, Preparing, each operation) +with nested per-host rows. Each node shows a spinner while running and a green +check / red cross (plus error details) when complete; verbose detail lines +(facts, command input/output) nest under the host nodes. + +Driven by ``pyinfra.api.state`` callbacks plus explicit phase context managers +for the synchronous setup steps. Only active on a TTY outside ``--json`` mode; +otherwise the CLI falls back to plain log lines. + +Concurrency note: nodes are mutated from many gevent greenlets without locks. +This is safe because greenlets are cooperative — mutations never yield midway — +and rendering happens from the Live refresh ticker between mutations. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING + +from rich.live import Live +from rich.progress_bar import ProgressBar +from rich.spinner import Spinner +from rich.table import Table +from rich.text import Text +from typing_extensions import override + +from pyinfra.api.renderable import OutputBlock +from pyinfra.api.state import BaseStateCallback + +from . import routing +from .console import console +from .renderables import to_renderable + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + + from pyinfra.api.host import Host + from pyinfra.api.state import State + +CHECK = Text("✓", style="bold green") +CROSS = Text("✗", style="bold red") +SKIP = Text("⤼", style="dim") # skipped (verbose only) +_SPINNER = "dots" +_REFRESH_PER_SECOND = 12.5 + +# Number of nested detail lines shown per node while it is still running (the +# full capture is rendered once the node completes / the frame persists). +RUNNING_DETAIL_LINES = 5 + + +class NodeStatus(str, Enum): + RUNNING = "running" + OK = "ok" + ERROR = "error" + SKIPPED = "skipped" + + +class NodeKind(str, Enum): + PHASE = "phase" + FILE = "file" + OPERATION = "operation" + HOST = "host" + + +# Progress-bar colours per node kind. Steps (phases + operations) are always +# cyan; files are a distinct grouping so they read as blue. +_BAR_STYLE = { + NodeKind.FILE: "blue", + NodeKind.OPERATION: "cyan", + NodeKind.PHASE: "cyan", +} + +# Detail-text / status colours matching the completion glyph. +_STATUS_STYLE = { + NodeStatus.OK: "green", + NodeStatus.ERROR: "red", + NodeStatus.SKIPPED: "dim", + NodeStatus.RUNNING: "dim", +} + + +def _grid() -> Table: + """The shared 4-column layout: glyph, label, progress bar/detail, count.""" + table = Table.grid(padding=(0, 1)) + table.add_column(width=1) + table.add_column() + table.add_column() + table.add_column() + return table + + +@dataclass +class DetailItem: + """A nested line under a host node: either plain text or a rich descriptor.""" + + text: str | None = None + descriptor: OutputBlock | None = None + is_error: bool = False + + @property + def is_renderable(self) -> bool: + return self.descriptor is not None + + +class Node: + """A single row in the progress tree.""" + + def __init__( + self, + label: str, + depth: int = 0, + total: int | None = None, + kind: NodeKind = NodeKind.PHASE, + ): + self.label = label + self.depth = depth + self.kind = kind + self.status = NodeStatus.RUNNING + self.detail: str | None = None + # Optional explicit style for the inline detail; falls back to the + # status colour when None. + self.detail_style: str | None = None + # Verbose per-host lines (fact loads, command input/output, diffs, ...) + # rendered nested under this node: list of (text, is_error, is_diff). + # A diff item's ``text`` is a whole diff block rendered via Syntax. + self.detail_lines: list[DetailItem] = [] + self.total = total # expected number of children (for the progress bar) + self.children: list[Node] = [] + self._spinner = Spinner(_SPINNER, style="cyan") + + def add(self, label: str, kind: NodeKind = NodeKind.HOST, total: int | None = None) -> Node: + child = Node(label, depth=self.depth + 1, total=total, kind=kind) + self.children.append(child) + return child + + def remove(self, child: Node) -> None: + if child in self.children: + self.children.remove(child) + + def add_detail_line(self, text: str, is_error: bool = False) -> None: + if text: + self.detail_lines.append(DetailItem(text=text, is_error=is_error)) + + def add_detail_renderable(self, descriptor: OutputBlock) -> None: + self.detail_lines.append(DetailItem(descriptor=descriptor)) + + def succeed(self, detail: str | None = None, detail_style: str | None = None) -> None: + self.status = NodeStatus.OK + self.detail = detail + self.detail_style = detail_style + + def fail(self, detail: str | None = None) -> None: + self.status = NodeStatus.ERROR + self.detail = detail + + def skip(self) -> None: + self.status = NodeStatus.SKIPPED + + @property + def is_parent(self) -> bool: + return self.total is not None or bool(self.children) + + @property + def completed(self) -> int: + # Skipped hosts are not part of ``total``, so exclude them here too. + return sum( + 1 for c in self.children if c.status not in (NodeStatus.RUNNING, NodeStatus.SKIPPED) + ) + + @property + def bar_total(self) -> int: + if self.total is not None: + return self.total + return len(self.children) + + def _glyph(self) -> Text | Spinner: + if self.status == NodeStatus.OK: + return CHECK + if self.status == NodeStatus.ERROR: + return CROSS + if self.status == NodeStatus.SKIPPED: + return SKIP + return self._spinner + + def render_rows(self, table: Table) -> None: + indent = " " * self.depth + if self.kind == NodeKind.HOST: + # Dim the "@connector/" prefix so the host name stands out. + base = "red" if self.status == NodeStatus.ERROR else "" + label = routing.host_label(self.label, base_style=base, prefix=indent) + else: + label = Text(f"{indent}{self.label}") + if self.status == NodeStatus.ERROR: + label.stylize("red") + elif self.kind == NodeKind.FILE: + label.stylize("bold blue") + elif self.kind in (NodeKind.PHASE, NodeKind.OPERATION): + # Match the label to the node's progress-bar colour. + label.stylize(_BAR_STYLE.get(self.kind, "")) + + if self.is_parent: + total = max(self.bar_total, 1) + completed = self.completed + if self.status == NodeStatus.ERROR: + bar_style = "red" + else: + bar_style = _BAR_STYLE.get(self.kind, "green") + bar = ProgressBar( + total=total, + completed=completed, + width=30, + finished_style=bar_style, + complete_style=bar_style, + ) + count = Text(f"{completed}/{self.bar_total}", style="dim") + table.add_row(self._glyph(), label, bar, count) + else: + row_detail = ( + Text( + self.detail, + style=self.detail_style or _STATUS_STYLE.get(self.status, "dim"), + ) + if self.detail + else Text("") + ) + table.add_row(self._glyph(), label, row_detail, Text("")) + + # Nested verbose detail lines: show the tail while running to keep the + # live region compact; render everything once the node has completed. + # Rich blocks (diffs, code blocks) are always kept; only plain text + # lines are tail-truncated while running. + if self.detail_lines: + lines = self.detail_lines + if self.status == NodeStatus.RUNNING: + plain = [item for item in lines if not item.is_renderable] + kept_plain = set(id(item) for item in plain[-RUNNING_DETAIL_LINES:]) + lines = [item for item in lines if item.is_renderable or id(item) in kept_plain] + detail_indent = " " * (self.depth + 1) + for item in lines: + if item.is_renderable: + assert item.descriptor is not None + table.add_row( + Text(""), + to_renderable(item.descriptor, indent=len(detail_indent)), + Text(""), + Text(""), + ) + else: + table.add_row( + Text(""), + Text( + f"{detail_indent}{item.text}", + style="red" if item.is_error else "dim", + ), + Text(""), + Text(""), + ) + + for child in self.children: + child.render_rows(table) + + +class DeployProgress(BaseStateCallback): + """ + State callback + live tree renderer. + + A single instance is created per run; it owns a Rich ``Live`` region and a + tree of :class:`Node` objects updated from state callbacks. Registered via + ``state.add_callback_handler``. Each phase — and each operation within the + Execute phase — gets its own live region so completed sections persist to + scrollback instead of being cropped when taller than the terminal. + + ``BaseStateCallback`` declares its hooks as ``@staticmethod`` but invokes + them via ``getattr(handler, name)``, so instance methods work fine at + runtime; the ``# type: ignore[override]`` markers below acknowledge the + intentional staticmethod→instance-method shape difference. + """ + + def __init__(self, state: State, verbose: bool = False): + self.state = state + self.verbose = verbose + self._roots: list[Node] = [] + self._op_nodes: dict[str, Node] = {} + self._op_host_nodes: dict[tuple[str, str], Node] = {} + self._file_nodes: dict[str, Node] = {} + self._error_hosts: list[tuple[str, Host]] = [] + self._connect_nodes: dict[str, Node] = {} + self._connect_root: Node | None = None + self._prepare_root: Node | None = None + self._prepare_nodes: dict[str, Node] = {} + # The most recent node for each host; verbose detail lines attach here. + self._active_host_node: dict[str, Node] = {} + self._paused = False + self._live = self._new_live() + + # Rendering + # + def _new_live(self) -> Live: + # ``get_renderable`` makes the Live pull (and build) the tree lazily at + # its own refresh rate instead of us re-rendering on every callback. + return Live( + get_renderable=self._render, + console=console, + refresh_per_second=_REFRESH_PER_SECOND, + transient=False, + ) + + def _render(self) -> Table: + table = _grid() + for root in self._roots: + root.render_rows(table) + return table + + @property + def is_active(self) -> bool: + """Whether routed host output can reach the display (running or paused).""" + return self._live.is_started or self._paused + + def add_step( + self, label: str, total: int | None = None, kind: NodeKind = NodeKind.PHASE + ) -> Node: + node = Node(label, total=total, kind=kind) + self._roots.append(node) + return node + + def __enter__(self) -> DeployProgress: + # Start a fresh live region for this phase so completed phases scroll + # up as static output and the new phase renders below. + self._roots = [] + self._file_nodes = {} + self._live = self._new_live() + self._live.start() + return self + + def __exit__(self, *exc: object) -> None: + self._finalize_pending_ops() + self._live.stop() + + def _rotate_region(self) -> None: + """Persist the current region to scrollback and start a fresh one.""" + self._finalize_pending_ops() + if self._live.is_started: + self._live.stop() + self._roots = [] + self._file_nodes = {} + self._live = self._new_live() + self._live.start() + + def _finalize_pending_ops(self) -> None: + """Mark still-running operation nodes succeed/fail from their children. + + ``operation_end`` finalizes op nodes on the normal path, but the + ``--serial`` / ``--no-wait`` paths never fire it, so op nodes created + lazily (see :meth:`_ensure_op_node`) would otherwise linger as running. + """ + for op_node in self._op_nodes.values(): + if op_node.status != NodeStatus.RUNNING: + continue + if any(c.status == NodeStatus.ERROR for c in op_node.children): + op_node.fail() + else: + op_node.succeed() + + def pause(self) -> bool: + """Clear and stop the live region (e.g. before an interactive prompt). + + Returns True if a live region was actually running (and should be + resumed afterwards with :meth:`resume`). + """ + if not self._live.is_started: + return False + # transient=True clears the region on stop instead of persisting it, + # so resume() doesn't render a duplicate frame below. + self._live.transient = True + self._live.stop() + self._paused = True + return True + + def resume(self) -> None: + """Restart the live region after :meth:`pause`, keeping the tree.""" + self._paused = False + self._live = self._new_live() + self._live.start() + + # Verbose detail routing + # + def add_host_detail(self, host_name: str, text: str, is_error: bool = False) -> None: + """Attach a routed log/echo line to the host's current tree node.""" + node = self._active_host_node.get(host_name) + if node is None: + return + # The per-host terminal status ("Success" / "No changes", optionally + # "... on retry N") is shown inline on the host node — green for a + # change, cyan for no change — mirroring the connect/prepare nodes, + # rather than as a dimmed detail line below. + if text.startswith("Success"): + node.succeed(text, detail_style="green") + return + if text.startswith("No changes"): + node.succeed(text, detail_style="cyan") + return + node.add_detail_line(text, is_error=is_error) + + def add_host_renderable(self, host_name: str, descriptor: OutputBlock) -> None: + """Attach a rich renderable block to the host's current tree node.""" + node = self._active_host_node.get(host_name) + if node is None: + return + node.add_detail_renderable(descriptor) + + # Host connect callbacks + # + @override + def host_before_connect(self, state: State, host: Host) -> None: # type: ignore[override] + if self._connect_root is None: + total = sum(1 for h in state.inventory if state.is_host_in_limit(h)) + self._connect_root = self.add_step("Connecting to hosts", total=total) + node = self._connect_root.add(host.name) + self._connect_nodes[host.name] = node + self._active_host_node[host.name] = node + + @override + def host_connect(self, state: State, host: Host) -> None: # type: ignore[override] + node = self._connect_nodes.get(host.name) + if node: + node.succeed("connected") + self._finish_connect_root() + + @override + def host_connect_error(self, state: State, host: Host, error) -> None: # type: ignore[override] + node = self._connect_nodes.get(host.name) + if node: + detail = str(error.args[0]) if getattr(error, "args", None) else str(error) + node.fail(detail) + self._finish_connect_root() + + def _finish_connect_root(self) -> None: + root = self._connect_root + if root is None: + return + children = root.children + # Hosts are added lazily as they start connecting: only conclude once + # every expected host has been added AND completed. + if len(children) != root.total: + return + if all(c.status != NodeStatus.RUNNING for c in children): + if any(c.status == NodeStatus.ERROR for c in children): + root.fail() + else: + root.succeed() + + # Prepare phase (driven directly by pyinfra_cli.util._parallel_load_hosts) + # + def prepare_start(self, name: str, hosts: Iterable[Host]) -> None: + """Start a "Preparing " phase with one child row per host.""" + hosts = list(hosts) + self._prepare_root = self.add_step(f"Preparing {name}", total=len(hosts)) + self._prepare_nodes = {} + for host in hosts: + node = self._prepare_root.add(host.name) + self._prepare_nodes[host.name] = node + self._active_host_node[host.name] = node + + def prepare_host_done(self, host: Host) -> None: + """Mark a host's prepare as complete (✗ if the host was failed).""" + node = self._prepare_nodes.get(host.name) + if node is None: + return + if host in self.state.failed_hosts: + errors = routing.get_host_errors().get(host.name) or [] + node.fail(errors[0] if errors else "failed") + else: + node.succeed("ready") + + def prepare_host_error(self, host: Host, error: BaseException) -> None: + node = self._prepare_nodes.get(host.name) + if node is None: + return + detail = str(error.args[0]) if getattr(error, "args", None) else str(error) + node.fail(detail) + + def prepare_end(self) -> None: + root = self._prepare_root + if root is None: + return + if any(c.status == NodeStatus.ERROR for c in root.children): + root.fail() + else: + root.succeed() + self._prepare_root = None + + # Operation callbacks + # + def _file_node(self, filename: str) -> Node: + """Get or create the parent node for a task/deploy file.""" + node = self._file_nodes.get(filename) + if node is None: + node = self.add_step(filename, kind=NodeKind.FILE) + self._file_nodes[filename] = node + return node + + @override + def operation_start(self, state: State, op_hash) -> None: # type: ignore[override] + # One live region per operation: persist the previous operation's + # subtree to scrollback so long deploys aren't cropped by the terminal + # height (Rich crops Live content taller than the screen). + if self._roots: + self._rotate_region() + + self._ensure_op_node(state, op_hash) + + def _ensure_op_node(self, state: State, op_hash) -> Node: + """Return the tree node for ``op_hash``, creating it if needed. + + ``operation_start`` normally creates it, but the ``--serial`` and + ``--no-wait`` execution paths never fire ``operation_start`` (they drive + hosts directly), so the node is created lazily on first host activity to + avoid dropping all host output under a TTY. + """ + op_node = self._op_nodes.get(op_hash) + if op_node is not None: + return op_node + + op_meta = state.get_op_meta(op_hash) + name = ", ".join(op_meta.names) if op_meta.names else "operation" + + # Operation names look like "path/to/file.py | Operation name". Nest the + # operation under a parent node for its file when present. + filename: str | None = None + if " | " in name: + filename, name = name.split(" | ", 1) + + # Count hosts that will actually run the op (failed hosts are excluded + # from the active set and never start). + total = sum(1 for host in state.inventory.get_active_hosts() if op_hash in state.ops[host]) + + if filename: + parent = self._file_node(filename) + op_node = parent.add(name, kind=NodeKind.OPERATION, total=total or None) + else: + op_node = self.add_step(name, total=total or None, kind=NodeKind.OPERATION) + + self._op_nodes[op_hash] = op_node + return op_node + + @override + def operation_host_start(self, state: State, host: Host, op_hash) -> None: # type: ignore[override] + # Lazily create the op node: --serial/--no-wait never fire + # operation_start, so without this the node (and all host output) would + # be dropped. + parent = self._ensure_op_node(state, op_hash) + node = parent.add(host.name, kind=NodeKind.HOST) + self._op_host_nodes[(op_hash, host.name)] = node + self._active_host_node[host.name] = node + + @override + def operation_host_skipped(self, state: State, host: Host, op_hash) -> None: # type: ignore[override] + # The host doesn't run this operation. By default drop the row entirely + # so it doesn't linger; in verbose mode keep it with a "skipped" glyph. + key = (op_hash, host.name) + node = self._op_host_nodes.get(key) + parent = self._op_nodes.get(op_hash) + if node is None or parent is None: + return + if self.verbose: + node.skip() + else: + # ``total`` already counts only hosts that run the op, so just drop + # the transient node created in operation_host_start. + parent.remove(node) + self._op_host_nodes.pop(key, None) + + @override + def operation_host_success( # type: ignore[override] + self, state: State, host: Host, op_hash, retry_count: int = 0 + ) -> None: + node = self._op_host_nodes.get((op_hash, host.name)) + if node: + # The inline status ("Success"/"No changes") is set from the routed + # status log line in add_host_detail (which fires first); mark the + # node OK without clobbering that detail. + node.succeed(node.detail, detail_style=node.detail_style) + + @override + def operation_host_error( # type: ignore[override] + self, state: State, host: Host, op_hash, retry_count: int = 0, max_retries: int = 0 + ) -> None: + node = self._op_host_nodes.get((op_hash, host.name)) + if node: + node.fail("failed") + self._error_hosts.append((op_hash, host)) + + @staticmethod + def _host_error_detail(state: State, host: Host, op_hash) -> str: + """Best-effort short error message from the operation's captured stderr.""" + try: + op_data = state.get_op_data_for_host(host, op_hash) + stderr = op_data.operation_meta.stderr_lines + except Exception: + stderr = [] + for line in stderr: + line = line.strip() + if line: + return line + return "failed" + + @override + def operation_end(self, state: State, op_hash) -> None: # type: ignore[override] + op_node = self._op_nodes.get(op_hash) + if op_node is None: + return + # Attach the real error message (from captured stderr) to failed hosts; + # operation_meta is only complete now, after all hosts have run. + for err_op_hash, host in self._error_hosts: + if err_op_hash != op_hash: + continue + node = self._op_host_nodes.get((op_hash, host.name)) + if node: + node.fail(self._host_error_detail(state, host, op_hash)) + if any(c.status == NodeStatus.ERROR for c in op_node.children): + op_node.fail() + else: + op_node.succeed() + + +def is_tree_active(json_output: bool) -> bool: + """The live tree is used on a TTY outside ``--json`` mode (all verbosity + levels — verbose detail lines nest under the host nodes).""" + if json_output: + return False + return console.is_terminal + + +@contextmanager +def step(progress: DeployProgress | None, label: str) -> Iterator[Node | None]: + """Run a synchronous step as a single spinner→check/cross row. + + Self-contained: renders its own short-lived ``Live`` region so it works + outside the phase live-regions owned by :class:`DeployProgress`. + """ + if progress is None: + yield None + return + + node = Node(label) + + def render() -> Table: + table = _grid() + node.render_rows(table) + return table + + with Live( + get_renderable=render, + console=console, + refresh_per_second=_REFRESH_PER_SECOND, + transient=False, + ): + try: + yield node + except Exception: + node.fail() + raise + else: + if node.status == NodeStatus.RUNNING: + node.succeed() diff --git a/src/pyinfra_cli/renderables.py b/src/pyinfra_cli/renderables.py new file mode 100644 index 000000000..ead37a203 --- /dev/null +++ b/src/pyinfra_cli/renderables.py @@ -0,0 +1,68 @@ +""" +Registry mapping rich-free output descriptors (:mod:`pyinfra.api.renderable`) to +Rich renderables. + +Operations/facts emit plain :class:`~pyinfra.api.renderable.OutputBlock` +descriptors via ``host.log_rich``; the CLI resolves each to a Rich renderable +here. New descriptor types register a renderer with +:func:`register_renderable`, so the routing/rendering pipeline stays generic. +""" + +from __future__ import annotations + +from collections.abc import Callable + +from rich.console import RenderableType +from rich.padding import Padding +from rich.syntax import Syntax +from rich.text import Text + +from pyinfra.api.renderable import CodeBlock, Diff, OutputBlock + +from .console import diff_renderable + +_RENDERERS: dict[type[OutputBlock], Callable[[OutputBlock], RenderableType]] = {} + + +def register_renderable( + cls: type[OutputBlock], + fn: Callable[[OutputBlock], RenderableType], +) -> None: + """Register the Rich renderer for a descriptor type.""" + _RENDERERS[cls] = fn + + +def _lookup(descriptor: OutputBlock) -> Callable[[OutputBlock], RenderableType] | None: + # Walk the MRO so a subclass without its own renderer falls back to a + # registered base renderer. + for klass in type(descriptor).__mro__: + fn = _RENDERERS.get(klass) # type: ignore[arg-type] + if fn is not None: + return fn + return None + + +def to_renderable(descriptor: OutputBlock, indent: int = 0) -> RenderableType: + """Turn a descriptor into a Rich renderable, indented if requested. + + Unknown descriptor types degrade to their ``repr`` rather than crashing. + """ + fn = _lookup(descriptor) + renderable: RenderableType = fn(descriptor) if fn is not None else Text(repr(descriptor)) + if indent: + return Padding(renderable, (0, 0, 0, indent)) + return renderable + + +# Built-in renderers (registered at import time). +register_renderable(Diff, lambda d: diff_renderable(d.text)) # type: ignore[attr-defined] +register_renderable( + CodeBlock, + lambda d: Syntax( + d.text, # type: ignore[attr-defined] + d.lexer, # type: ignore[attr-defined] + background_color="default", + theme="ansi_dark", + word_wrap=True, + ), +) diff --git a/src/pyinfra_cli/routing.py b/src/pyinfra_cli/routing.py new file mode 100644 index 000000000..92e02c276 --- /dev/null +++ b/src/pyinfra_cli/routing.py @@ -0,0 +1,139 @@ +""" +Host attribution & routing of log/echo messages. + +When the live progress tree is active, per-host log and echo lines are routed +into the host's tree node instead of streaming to the console (which would +corrupt the live region). Warnings/errors are also recorded per host so +failure prompts can display *which* hosts failed and why. + +Attribution uses ``pyinfra.context.ctx_host`` when set in the calling greenlet +and falls back to parsing the rendered ``host.print_prefix`` at the start of +the message (command-output reader greenlets and connect-phase greenlets do +not inherit the host context, but their lines carry the prefix). +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +from rich.text import Text + +from pyinfra.context import ctx_host + +if TYPE_CHECKING: + from collections.abc import Iterable + + from .progress import DeployProgress + +ANSI_RE = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]") + +# Matches a leading connector prefix, e.g. "@docker/" or "@fake/". +CONNECTOR_PREFIX_RE = re.compile(r"^(@[^/]+/)(.*)$") + + +def host_label(name: str, base_style: str = "", prefix: str = "") -> Text: + """ + Render a host name as Rich ``Text`` with any ``@connector/`` prefix dimmed. + + ``prefix`` is prepended verbatim (e.g. tree indentation). + """ + text = Text(prefix) + match = CONNECTOR_PREFIX_RE.match(name) + if match: + connector, rest = match.groups() + text.append(connector, style="dim") + text.append(rest, style=base_style) + else: + text.append(name, style=base_style) + return text + + +_tree: DeployProgress | None = None +_host_names: list[str] = [] # sorted longest-first for prefix matching +_host_names_set: set[str] = set() +_has_whitespace_names = False +_host_errors: dict[str, list[str]] = {} + + +def strip_ansi(text: str) -> str: + return ANSI_RE.sub("", text) + + +def set_tree(tree: DeployProgress | None) -> None: + """Register (or clear) the active live tree renderer.""" + global _tree + _tree = tree + + +def get_tree() -> DeployProgress | None: + return _tree + + +def set_host_names(names: Iterable[str]) -> None: + """Register the inventory host names used for prefix attribution.""" + global _host_names, _host_names_set, _has_whitespace_names + _host_names = sorted(names, key=len, reverse=True) + _host_names_set = set(_host_names) + _has_whitespace_names = any(" " in name for name in _host_names) + + +def reset_host_errors() -> None: + _host_errors.clear() + + +def record_host_error(host_name: str, message: str) -> None: + _host_errors.setdefault(host_name, []).append(message) + + +def get_host_errors() -> dict[str, list[str]]: + return _host_errors + + +def split_host_prefix(plain: str) -> tuple[str | None, str]: + """ + Split a plain (ANSI-stripped) message into ``(host_name, rest)`` when it + starts with a known host prefix, else ``(None, message)``. + + Supports both the legacy ``[hostname]`` bracketed form and the plain + ``hostname`` form. The prefix is always ``name`` + padding + space, so the + common case is an O(1) lookup of the first whitespace-delimited token; the + linear prefix scan only remains for host names containing whitespace. + """ + s = plain.lstrip() + + if s.startswith("["): + end = s.find("]") + if end > 0 and s[1:end] in _host_names_set: + return s[1:end], s[end + 1 :].lstrip() + + token = s.split(maxsplit=1)[0] if s else "" + if token in _host_names_set: + return token, s[len(token) :].lstrip() + + if _has_whitespace_names: + for name in _host_names: + if s.startswith(name): + return name, s[len(name) :].lstrip() + + return None, plain.strip() + + +def attribute_host(message: str) -> tuple[str | None, str]: + """ + Attribute a rendered log/echo ``message`` to a host. + + Returns ``(host_name | None, plain_detail_text)`` — the detail text is + ANSI-stripped with any host prefix removed. + """ + plain = strip_ansi(message) + + name, rest = split_host_prefix(plain) + if name is not None: + return name, rest + + if ctx_host.isset(): + host = ctx_host.get() + return host.name, plain.strip() + + return None, plain.strip() diff --git a/src/pyinfra_cli/util.py b/src/pyinfra_cli/util.py index cdaefd306..a493fff7b 100644 --- a/src/pyinfra_cli/util.py +++ b/src/pyinfra_cli/util.py @@ -215,6 +215,14 @@ def try_import_module_attribute(path, prefix=None, raise_for_none=True): def _parallel_load_hosts(state: State, callback: Callable, name: str): + from .routing import get_tree + + tree = get_tree() + hosts = list(state.inventory.get_active_hosts()) + + if tree is not None: + tree.prepare_start(name, hosts) + def load_file(local_host): try: with ctx_config.use(state.config.copy()): @@ -226,18 +234,30 @@ def load_file(local_host): except Exception as e: return e - greenlet_to_host = { - state.pool.spawn(load_file, host): host for host in state.inventory.get_active_hosts() - } + greenlet_to_host = {state.pool.spawn(load_file, host): host for host in hosts} + + # Wait for *all* hosts to finish evaluating before raising any error, so + # host status/output isn't interleaved with error handling or prompts. + errors: list[Exception] = [] with progress_spinner(greenlet_to_host.values()) as progress: for greenlet in gevent.iwait(greenlet_to_host.keys()): host = greenlet_to_host[greenlet] result = greenlet.get() if isinstance(result, Exception): - raise result + errors.append(result) + if tree is not None: + tree.prepare_host_error(host, result) + elif tree is not None: + tree.prepare_host_done(host) progress(host) + if tree is not None: + tree.prepare_end() + + if errors: + raise errors[0] + def load_deploy_file(state: State, filename): state.current_deploy_filename = filename diff --git a/tests/end-to-end/conftest.py b/tests/end-to-end/conftest.py index ee6ce942a..9ac22cf27 100644 --- a/tests/end-to-end/conftest.py +++ b/tests/end-to-end/conftest.py @@ -31,7 +31,7 @@ def run(command, cwd=None, expected_exit_code=0): @staticmethod def run_check_output(command, expected_lines=None, **kwargs): if expected_lines is None: - expected_lines = ["Connected", "Starting operation", "Errors: 0"] + expected_lines = ["Connected", "Starting operation", "Grand total"] _, stderr = Helpers.run(command, **kwargs) diff --git a/tests/end-to-end/test_e2e_local.py b/tests/end-to-end/test_e2e_local.py index b2ef5feb8..8a381725c 100644 --- a/tests/end-to-end/test_e2e_local.py +++ b/tests/end-to-end/test_e2e_local.py @@ -22,13 +22,13 @@ def temp_dir(): def test_int_local_file_no_changes(helpers, temp_dir): helpers.run_check_output( # first run = create the file "pyinfra -y -v @local files.file _testfile", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=temp_dir, ) helpers.run_check_output( # second run = no changes "pyinfra -y -v @local files.file _testfile", - expected_lines=["@local] No changes"], + expected_lines=[r"@local\s+No changes"], cwd=temp_dir, ) @@ -38,25 +38,25 @@ def test_int_local_file_no_changes(helpers, temp_dir): def test_int_local_directory_no_changes(helpers, temp_dir): helpers.run_check_output( # first run = create the directory "pyinfra -y -v @local files.directory _testdir", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=temp_dir, ) helpers.run_check_output( # second run = no changes "pyinfra -y -v @local files.directory _testdir", - expected_lines=["@local] No changes"], + expected_lines=[r"@local\s+No changes"], cwd=temp_dir, ) helpers.run_check_output( # third run (remove) = remove directory "pyinfra -y -v @local files.directory _testdir present=False", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=temp_dir, ) helpers.run_check_output( # fourth run (remove) = no chances "pyinfra -y -v @local files.directory _testdir present=False", - expected_lines=["@local] No changes"], + expected_lines=[r"@local\s+No changes"], cwd=temp_dir, ) @@ -66,13 +66,13 @@ def test_int_local_directory_no_changes(helpers, temp_dir): def test_int_local_link_no_changes(helpers, temp_dir): helpers.run_check_output( # first run = create the link "pyinfra -y -v @local files.link _testlink target=_testfile", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=temp_dir, ) helpers.run_check_output( # second run = no changes "pyinfra -y -v @local files.link _testlink target=_testfile", - expected_lines=["@local] No changes"], + expected_lines=[r"@local\s+No changes"], cwd=temp_dir, ) @@ -82,25 +82,25 @@ def test_int_local_link_no_changes(helpers, temp_dir): def test_int_local_line_no_changes(helpers, temp_dir): helpers.run_check_output( # first run = create the line "pyinfra -y -v @local files.line _testfile someline", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=temp_dir, ) helpers.run_check_output( # second run = no changes "pyinfra -y -v @local files.line _testfile someline", - expected_lines=["@local] No changes"], + expected_lines=[r"@local\s+No changes"], cwd=temp_dir, ) helpers.run_check_output( # replace the line "pyinfra -y -v @local files.line _testfile someline replace=anotherline", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=temp_dir, ) helpers.run_check_output( # second run replace the line = no changes "pyinfra -y -v @local files.line _testfile someline replace=anotherline", - expected_lines=["@local] No changes"], + expected_lines=[r"@local\s+No changes"], cwd=temp_dir, ) @@ -113,7 +113,7 @@ def test_int_local_line_ensure_newline_true(helpers, tmp_path): path.write_bytes(b"hello world") helpers.run_check_output( "pyinfra -y -v @local files.line _testfile someline ensure_newline=true", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=tmp_path, ) assert path.read_bytes() == b"hello world\nsomeline\n" @@ -121,7 +121,7 @@ def test_int_local_line_ensure_newline_true(helpers, tmp_path): path.write_bytes(b"hello world\n") helpers.run_check_output( "pyinfra -y -v @local files.line _testfile someline ensure_newline=true", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=tmp_path, ) assert path.read_bytes() == b"hello world\nsomeline\n" @@ -135,7 +135,7 @@ def test_int_local_line_ensure_newline_false(helpers, tmp_path): path.write_bytes(b"hello world") helpers.run_check_output( "pyinfra -y -v @local files.line _testfile someline ensure_newline=false", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=tmp_path, ) assert path.read_bytes() == b"hello worldsomeline\n" @@ -143,7 +143,7 @@ def test_int_local_line_ensure_newline_false(helpers, tmp_path): path.write_bytes(b"hello world\n") helpers.run_check_output( "pyinfra -y -v @local files.line _testfile someline ensure_newline=false", - expected_lines=["@local] Success"], + expected_lines=[r"@local\s+Success"], cwd=tmp_path, ) assert path.read_bytes() == b"hello world\nsomeline\n" diff --git a/tests/end-to-end/test_e2e_ssh.py b/tests/end-to-end/test_e2e_ssh.py index 90a050f9b..b09cce204 100644 --- a/tests/end-to-end/test_e2e_ssh.py +++ b/tests/end-to-end/test_e2e_ssh.py @@ -48,11 +48,11 @@ def run_docker_ssh_server(helpers): def test_e2e_ssh_sudo_password(helpers): helpers.run_check_output( f"{PYINFRA_COMMAND} server.shell echo _sudo=True _sudo_password=password", - expected_lines=["localhost] Success"], + expected_lines=[r"localhost\s+Success"], ) helpers.run_check_output( f"{PYINFRA_COMMAND} server.shell echo _sudo=True _sudo_password=wrongpassword", - expected_lines=["localhost] Error"], + expected_lines=[r"localhost\s+Error"], expected_exit_code=1, ) @@ -62,12 +62,12 @@ def test_e2e_ssh_sudo_password(helpers): def test_int_local_file_no_changes(helpers): helpers.run_check_output( # first run = create the file f"{PYINFRA_COMMAND} files.file _testfile", - expected_lines=["localhost] Success"], + expected_lines=[r"localhost\s+Success"], ) helpers.run_check_output( # second run = no changes f"{PYINFRA_COMMAND} files.file _testfile", - expected_lines=["localhost] No changes"], + expected_lines=[r"localhost\s+No changes"], ) @@ -76,22 +76,22 @@ def test_int_local_file_no_changes(helpers): def test_int_local_directory_no_changes(helpers): helpers.run_check_output( # first run = create the directory f"{PYINFRA_COMMAND} files.directory _testdir", - expected_lines=["localhost] Success"], + expected_lines=[r"localhost\s+Success"], ) helpers.run_check_output( # second run = no changes f"{PYINFRA_COMMAND} files.directory _testdir", - expected_lines=["localhost] No changes"], + expected_lines=[r"localhost\s+No changes"], ) helpers.run_check_output( # third run (remove) = remove directory f"{PYINFRA_COMMAND} files.directory _testdir present=False", - expected_lines=["localhost] Success"], + expected_lines=[r"localhost\s+Success"], ) helpers.run_check_output( # fourth run (remove) = no chances f"{PYINFRA_COMMAND} files.directory _testdir present=False", - expected_lines=["localhost] No changes"], + expected_lines=[r"localhost\s+No changes"], ) @@ -100,12 +100,12 @@ def test_int_local_directory_no_changes(helpers): def test_int_local_link_no_changes(helpers): helpers.run_check_output( # first run = create the link f"{PYINFRA_COMMAND} files.link _testlink target=_testfile", - expected_lines=["localhost] Success"], + expected_lines=[r"localhost\s+Success"], ) helpers.run_check_output( # second run = no changes f"{PYINFRA_COMMAND} files.link _testlink target=_testfile", - expected_lines=["localhost] No changes"], + expected_lines=[r"localhost\s+No changes"], ) @@ -114,20 +114,20 @@ def test_int_local_link_no_changes(helpers): def test_int_local_line_no_changes(helpers): helpers.run_check_output( # first run = create the line f"{PYINFRA_COMMAND} files.line _testfile someline", - expected_lines=["localhost] Success"], + expected_lines=[r"localhost\s+Success"], ) helpers.run_check_output( # second run = no changes f"{PYINFRA_COMMAND} files.line _testfile someline", - expected_lines=["localhost] No changes"], + expected_lines=[r"localhost\s+No changes"], ) helpers.run_check_output( # replace the line f"{PYINFRA_COMMAND} files.line _testfile someline replace=anotherline", - expected_lines=["localhost] Success"], + expected_lines=[r"localhost\s+Success"], ) helpers.run_check_output( # second run replace the line = no changes f"{PYINFRA_COMMAND} files.line _testfile someline replace=anotherline", - expected_lines=["localhost] No changes"], + expected_lines=[r"localhost\s+No changes"], ) diff --git a/tests/test_cli/test_cli_log.py b/tests/test_cli/test_cli_log.py new file mode 100644 index 000000000..27ecbc771 --- /dev/null +++ b/tests/test_cli/test_cli_log.py @@ -0,0 +1,70 @@ +"""Unit tests for the CLI log handler's rich-descriptor routing.""" + +import logging +from unittest.mock import MagicMock, patch + +from pyinfra.api.renderable import Diff + +from pyinfra_cli.log import LogHandler + + +def _rich_record(descriptor, host="@fake/web-1") -> logging.LogRecord: + record = logging.LogRecord( + name="pyinfra", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="", + args=(), + exc_info=None, + ) + record.pyinfra_rich = descriptor + record.pyinfra_host = host + return record + + +DIFF = Diff("@@ -1,1 +1,1 @@\n- old line\n+ new line") + + +class TestRichRouting: + def test_descriptor_routed_to_tree(self): + tree = MagicMock() + tree.is_active = True + handler = LogHandler() + + with patch("pyinfra_cli.log.routing.get_tree", return_value=tree): + handler.emit(_rich_record(DIFF)) + + tree.add_host_renderable.assert_called_once() + host_name, descriptor = tree.add_host_renderable.call_args.args + assert host_name == "@fake/web-1" + # The descriptor is passed intact (not stringified / re-parsed). + assert descriptor is DIFF + + def test_descriptor_printed_when_no_tree(self): + handler = LogHandler() + + with ( + patch("pyinfra_cli.log.routing.get_tree", return_value=None), + patch("pyinfra_cli.log.console") as console, + ): + handler.emit(_rich_record(DIFF)) + + console.print.assert_called_once() + # Rendered via the registry (a Syntax/Padding renderable), not a string. + (renderable,) = console.print.call_args.args + assert not isinstance(renderable, str) + + def test_descriptor_printed_when_tree_inactive(self): + tree = MagicMock() + tree.is_active = False + handler = LogHandler() + + with ( + patch("pyinfra_cli.log.routing.get_tree", return_value=tree), + patch("pyinfra_cli.log.console") as console, + ): + handler.emit(_rich_record(DIFF)) + + tree.add_host_renderable.assert_not_called() + console.print.assert_called_once() diff --git a/tests/test_cli/test_cli_progress.py b/tests/test_cli/test_cli_progress.py new file mode 100644 index 000000000..5e25afeae --- /dev/null +++ b/tests/test_cli/test_cli_progress.py @@ -0,0 +1,221 @@ +"""Unit tests for the live progress tree routing and the TTY output-loss fixes. + +The live tree is only active on a TTY, which pytest cannot easily provide, so +these exercise the routing logic directly rather than through a subprocess. +""" + +from unittest.mock import MagicMock + +from pyinfra_cli.log import _is_lifecycle_noise +from pyinfra.api.renderable import Diff + +from pyinfra_cli.progress import DeployProgress, DetailItem, NodeKind, NodeStatus + + +def _make_progress(verbose: bool = False) -> DeployProgress: + # DeployProgress only needs `state` for the op-node helpers, which we drive + # with a mock below, so a bare mock state is enough here. + return DeployProgress(MagicMock(), verbose=verbose) + + +class TestLifecycleNoiseGate: + """Fix B/C: only lifecycle noise is verbose-gated; deploy output is not.""" + + def test_lifecycle_messages_are_noise(self): + assert _is_lifecycle_noise("Connected") + assert _is_lifecycle_noise("Ready: apt.packages") + assert _is_lifecycle_noise("Disconnected") + assert _is_lifecycle_noise("noop: user already exists") + + def test_deploy_output_is_not_noise(self): + # Real deploy output / diffs must always route to the host node. + assert not _is_lifecycle_noise("Will modify /etc/hosts") + assert not _is_lifecycle_noise("- old line") + assert not _is_lifecycle_noise("+ new line") + assert not _is_lifecycle_noise("some command output") + + +class TestSerialNoWaitLazyNodes: + """Fix A: --serial/--no-wait never fire operation_start, so the op node must + be created lazily on first host activity or all host output is dropped.""" + + def _mock_state_for_op(self, op_hash: str, host): + state = MagicMock() + op_meta = MagicMock() + op_meta.names = ["server.shell (echo hi)"] + state.get_op_meta.return_value = op_meta + state.inventory.get_active_hosts.return_value = [host] + state.ops = {host: {op_hash: object()}} + return state + + def test_operation_host_start_creates_op_node_without_operation_start(self): + progress = _make_progress() + host = MagicMock() + host.name = "@fake/web-1" + op_hash = "abc123" + state = self._mock_state_for_op(op_hash, host) + + # operation_start is intentionally NOT called (mirrors --serial/--no-wait) + progress.operation_host_start(state, host, op_hash) + + # The op node was lazily created, and the host node registered so + # add_host_detail has somewhere to attach output. + assert op_hash in progress._op_nodes + assert progress._active_host_node.get(host.name) is not None + + # Host-attributed output is now retained rather than dropped. + progress.add_host_detail(host.name, "hello from serial") + node = progress._active_host_node[host.name] + assert DetailItem(text="hello from serial") in node.detail_lines + + def test_finalize_pending_ops_marks_running_op_success(self): + progress = _make_progress() + host = MagicMock() + host.name = "@fake/web-1" + op_hash = "abc123" + state = self._mock_state_for_op(op_hash, host) + + progress.operation_host_start(state, host, op_hash) + progress.operation_host_success(state, host, op_hash) + # operation_end never fires under --serial/--no-wait; the op node lingers + # as RUNNING until finalised. + assert progress._op_nodes[op_hash].status == NodeStatus.RUNNING + + progress._finalize_pending_ops() + assert progress._op_nodes[op_hash].status == NodeStatus.OK + + def test_finalize_pending_ops_marks_failed_op_error(self): + progress = _make_progress() + host = MagicMock() + host.name = "@fake/web-1" + op_hash = "abc123" + state = self._mock_state_for_op(op_hash, host) + + progress.operation_host_start(state, host, op_hash) + progress.operation_host_error(state, host, op_hash) + + progress._finalize_pending_ops() + assert progress._op_nodes[op_hash].status == NodeStatus.ERROR + + +class TestAddHostDetailDropsWhenNoNode: + """add_host_detail is a safe no-op when a host has no active node (e.g. a + line arrives before any phase started); it must never raise.""" + + def test_no_node_is_noop(self): + progress = _make_progress() + # Should not raise even though no node exists for this host. + progress.add_host_detail("@fake/unknown", "orphan line") + assert progress._active_host_node.get("@fake/unknown") is None + + +class TestNodeKindEnum: + def test_kinds_exist(self): + assert NodeKind.OPERATION + assert NodeKind.HOST + + +class TestInlineHostStatus: + """The per-host terminal status shows inline on the host node (green + Success / cyan No changes), not as a dimmed detail line below.""" + + def _host_node(self): + progress = _make_progress() + host = MagicMock() + host.name = "@fake/web-1" + op_hash = "abc123" + op_meta = MagicMock() + op_meta.names = ["server.shell"] + state = MagicMock() + state.get_op_meta.return_value = op_meta + state.inventory.get_active_hosts.return_value = [host] + state.ops = {host: {op_hash: object()}} + progress.operation_host_start(state, host, op_hash) + return progress, progress._active_host_node[host.name] + + def test_success_is_inline_and_green(self): + progress, node = self._host_node() + progress.add_host_detail("@fake/web-1", "Success") + assert node.status == NodeStatus.OK + assert node.detail == "Success" + assert node.detail_style == "green" + # Not duplicated as a dimmed detail line. + assert node.detail_lines == [] + + def test_no_changes_is_inline_and_cyan(self): + progress, node = self._host_node() + progress.add_host_detail("@fake/web-1", "No changes") + assert node.detail == "No changes" + assert node.detail_style == "cyan" + assert node.detail_lines == [] + + def test_other_output_still_a_detail_line(self): + progress, node = self._host_node() + progress.add_host_detail("@fake/web-1", "some command output") + assert node.detail_lines == [DetailItem(text="some command output")] + + def test_operation_host_success_keeps_inline_detail(self): + progress, node = self._host_node() + # Status line arrives first, then the success callback fires. + progress.add_host_detail("@fake/web-1", "Success") + state = MagicMock() + host = MagicMock() + host.name = "@fake/web-1" + progress.operation_host_success(state, host, "abc123") + # The callback must not clobber the inline detail. + assert node.detail == "Success" + assert node.detail_style == "green" + + +class TestHostRenderable: + """Rich descriptor blocks (diffs, code blocks) are stored as renderable + detail items and rendered via the registry; they are exempt from the + running tail-truncation.""" + + def _host_node(self): + progress = _make_progress() + host = MagicMock() + host.name = "@fake/web-1" + op_hash = "abc123" + op_meta = MagicMock() + op_meta.names = ["files.put"] + state = MagicMock() + state.get_op_meta.return_value = op_meta + state.inventory.get_active_hosts.return_value = [host] + state.ops = {host: {op_hash: object()}} + progress.operation_host_start(state, host, op_hash) + return progress, progress._active_host_node[host.name] + + def test_add_host_renderable_appends_item(self): + progress, node = self._host_node() + diff = Diff("@@ -1,1 +1,1 @@\n- old\n+ new") + progress.add_host_renderable("@fake/web-1", diff) + assert node.detail_lines == [DetailItem(descriptor=diff)] + + def test_add_host_renderable_no_node_is_noop(self): + progress = _make_progress() + # No active node for this host -> must not raise. + progress.add_host_renderable("@fake/unknown", Diff("@@ -1 +1 @@\n- a\n+ b")) + + def test_renderable_item_survives_running_truncation(self): + progress, node = self._host_node() + # Many plain lines + one renderable; while RUNNING plain lines are + # tail-trimmed but the renderable is always kept. + for i in range(10): + node.add_detail_line(f"line {i}") + node.add_detail_renderable(Diff("@@ -1 +1 @@\n- a\n+ b")) + assert node.status == NodeStatus.RUNNING + + table = MagicMock() + rendered = [] + table.add_row.side_effect = lambda *cols: rendered.append(cols) + node.render_rows(table) + + # The renderable block is rendered (Padding-wrapped) even though most + # plain lines were tail-truncated. + from rich.padding import Padding + + assert any(isinstance(cols[1], Padding) for cols in rendered) + # Plain lines were tail-truncated (at most RUNNING_DETAIL_LINES kept). + plain_kept = sum(1 for cols in rendered if "line " in str(cols)) + assert plain_kept <= 5 diff --git a/tests/test_cli/test_cli_renderables.py b/tests/test_cli/test_cli_renderables.py new file mode 100644 index 000000000..df1248bc4 --- /dev/null +++ b/tests/test_cli/test_cli_renderables.py @@ -0,0 +1,61 @@ +"""Unit tests for the CLI descriptor -> Rich renderable registry.""" + +from dataclasses import dataclass + +from rich.padding import Padding +from rich.syntax import Syntax +from rich.text import Text + +from pyinfra.api.renderable import CodeBlock, Diff, OutputBlock + +from pyinfra_cli import renderables +from pyinfra_cli.renderables import register_renderable, to_renderable + + +class TestBuiltinRenderers: + def test_diff_renders_as_syntax(self): + r = to_renderable(Diff("@@ -1 +1 @@\n- a\n+ b")) + assert isinstance(r, Syntax) + assert r.lexer.name.lower() == "diff" + + def test_code_block_uses_its_lexer(self): + r = to_renderable(CodeBlock("SELECT 1;", "sql")) + assert isinstance(r, Syntax) + assert r.lexer.name.lower() == "sql" + + def test_indent_wraps_in_padding(self): + r = to_renderable(Diff("@@ -1 +1 @@\n- a\n+ b"), indent=4) + assert isinstance(r, Padding) + + +class TestRegistry: + def test_subclass_falls_back_to_base_renderer(self): + @dataclass(frozen=True) + class SpecialDiff(Diff): + pass + + # No renderer registered for SpecialDiff -> resolves to Diff's via MRO. + r = to_renderable(SpecialDiff("@@ -1 +1 @@\n- a\n+ b")) + assert isinstance(r, Syntax) + assert r.lexer.name.lower() == "diff" + + def test_unknown_descriptor_degrades_to_repr(self): + @dataclass(frozen=True) + class Unknown(OutputBlock): + value: int = 1 + + r = to_renderable(Unknown()) + assert isinstance(r, Text) + assert "Unknown" in str(r) + + def test_register_renderable(self): + @dataclass(frozen=True) + class Custom(OutputBlock): + body: str = "" + + marker = Text("custom!") + register_renderable(Custom, lambda d: marker) + try: + assert to_renderable(Custom("x")) is marker + finally: + renderables._RENDERERS.pop(Custom, None) diff --git a/tests/test_operations_utils.py b/tests/test_operations_utils.py index 8523c2537..537c9b160 100644 --- a/tests/test_operations_utils.py +++ b/tests/test_operations_utils.py @@ -1,14 +1,49 @@ from unittest import TestCase -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest from pyinfra.facts.util.packages import PackageInfo, PackageStatus from pyinfra.operations.util.docker import parse_image_reference, parse_registry -from pyinfra.operations.util.files import ensure_mode_int, unix_path_join +from pyinfra.operations.util.files import ( + ensure_mode_int, + generate_color_diff, + generate_diff, + unix_path_join, +) from pyinfra.operations.util.packaging import ensure_packages +class TestGenerateDiff(TestCase): + def test_plain_diff_markers(self): + lines = list(generate_diff(["a\n", "b\n", "c\n"], ["a\n", "X\n", "c\n", "d\n"])) + text = "\n".join(lines) + # Plain markers, no ANSI escapes. + assert "\x1b" not in text + assert any(line.startswith("@@ ") for line in lines) + assert "- b" in lines + assert "+ X" in lines + assert "+ d" in lines + assert " a" in lines # context line + + def test_empty_diff(self): + assert list(generate_diff(["a\n"], ["a\n"])) == [] + + +class TestGenerateColorDiff(TestCase): + def test_color_diff_wraps_changes_with_format_text(self): + # ``format_text`` is terminal-gated (plain when not a TTY), so assert the + # colour wrapping is applied to -/+ lines rather than checking raw ANSI. + with patch("pyinfra.operations.util.files.format_text") as fake_format: + fake_format.side_effect = lambda text, fg: f"<{fg}>{text}" + lines = list(generate_color_diff(["a\n", "b\n"], ["a\n", "X\n"])) + assert "- b" in lines + assert "+ X" in lines + # Hunk header and context stay plain. + assert "@@ -1,2 +1,2 @@" in lines + assert " a" in lines + + class TestUnixPathJoin(TestCase): def test_simple_path(self): assert unix_path_join("home", "pyinfra") == "home/pyinfra"