From a9a68ec22712a6535ed4386fa46b07c7aa990c37 Mon Sep 17 00:00:00 2001 From: "Axel H." Date: Sat, 4 Jul 2026 18:10:39 +0200 Subject: [PATCH 1/3] refactor(cli): migrate from Click to Cyclopts + Rich Replace Click with Cyclopts for argument parsing/help and Rich for all output. - pyinfra.api.output gains a shared Rich Console; format_text/echo are Rich-backed adapters preserving the legacy click.style signature - pyinfra_cli.console holds the stderr (human) and stdout (JSON) consoles - cli.py uses a single Cyclopts @app.default preserving the exact CLI syntax (inventory + variadic operations, faked subcommand dispatch, exec -- passthrough, -v counting, --user/--port aliases, PYINFRA_YES) - exceptions.py drops click.ClickException for Rich-rendered CliException - CLI test harness swaps click.testing.CliRunner for a Cyclopts-based run_cli - Shell completion installed via `pyinfra --install-completion` (Cyclopts); the old Click completion scripts are removed and docs/CHANGELOG updated stdout stays byte-for-byte JSON; all human output remains on stderr. --- docs/cli.md | 19 +- pyproject.toml | 4 +- scripts/pyinfra-complete.sh | 21 -- scripts/pyinfra-complete.zsh | 28 -- src/pyinfra/api/output.py | 51 ++- src/pyinfra/progress.py | 141 ++----- src/pyinfra_cli/cli.py | 523 +++++++++++++------------- src/pyinfra_cli/console.py | 70 ++++ src/pyinfra_cli/exceptions.py | 133 ++++--- src/pyinfra_cli/log.py | 17 +- src/pyinfra_cli/main.py | 23 +- src/pyinfra_cli/prints.py | 77 ++-- src/pyinfra_cli/util.py | 4 +- src/pyinfra_cli/virtualenv.py | 6 +- tests/test_cli/test_cli.py | 53 +++ tests/test_cli/test_cli_exceptions.py | 8 +- tests/test_cli/util.py | 65 +++- uv.lock | 101 ++++- 18 files changed, 771 insertions(+), 573 deletions(-) delete mode 100644 scripts/pyinfra-complete.sh delete mode 100644 scripts/pyinfra-complete.zsh create mode 100644 src/pyinfra_cli/console.py diff --git a/docs/cli.md b/docs/cli.md index 553f641e9..828201d27 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -8,7 +8,7 @@ pyinfra is an extremely powerful tool for ad-hoc execution and management of rem As described in the [getting started page](./getting-started), pyinfra needs an **inventory** and some **operations**. These are used with the CLI as below: ```sh -Usage: pyinfra [OPTIONS] INVENTORY OPERATIONS... +Usage: pyinfra [OPTIONS] INVENTORY [OPERATIONS...] # INVENTORY @@ -222,14 +222,15 @@ Interactive prompts from deploy code (e.g. `input(...)`) will still block, so av ## Shell Autocompletion -Add the following to your `~/.bash_profile` or `~/.profile` files: +pyinfra's CLI is built on [Cyclopts](https://cyclopts.readthedocs.io), which can install +shell completion for you: -+ **bash** `source scripts/pyinfra-complete.sh`. -+ **zsh** `source scripts/pyinfra-complete.zsh`. - -These files were generated using these commands: +```sh +# Install completion for your current shell (bash, zsh, fish, ...) +pyinfra --install-completion +# Or target a specific shell +pyinfra --install-completion --shell zsh ``` -env _PYINFRA_COMPLETE=bash_source pyinfra > pyinfra-complete.sh -env _PYINFRA_COMPLETE=zsh_source pyinfra > pyinfra-complete.zsh -``` + +After installing you may need to restart your shell or source its configuration file. diff --git a/pyproject.toml b/pyproject.toml index 47519fbc4..f670a1729 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,8 @@ requires-python = ">=3.10,<4.0" dependencies = [ "gevent>=1.5", "paramiko>=2.11,<5", # 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", @@ -49,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/scripts/pyinfra-complete.sh b/scripts/pyinfra-complete.sh deleted file mode 100644 index d073dbf38..000000000 --- a/scripts/pyinfra-complete.sh +++ /dev/null @@ -1,21 +0,0 @@ -_pyinfra_completion() { - local IFS=$' -' - COMPREPLY=( $( env COMP_WORDS="${COMP_WORDS[*]}" \ - COMP_CWORD=$COMP_CWORD \ - _PYINFRA_COMPLETE=complete $1 ) ) - return 0 -} - -_pyinfra_completionetup() { - local COMPLETION_OPTIONS="" - local BASH_VERSION_ARR=(${BASH_VERSION//./ }) - # Only BASH version 4.4 and later have the nosort option. - if [ ${BASH_VERSION_ARR[0]} -gt 4 ] || ([ ${BASH_VERSION_ARR[0]} -eq 4 ] && [ ${BASH_VERSION_ARR[1]} -ge 4 ]); then - COMPLETION_OPTIONS="-o nosort" - fi - - complete $COMPLETION_OPTIONS -F _pyinfra_completion pyinfra -} - -_pyinfra_completionetup; diff --git a/scripts/pyinfra-complete.zsh b/scripts/pyinfra-complete.zsh deleted file mode 100644 index 8f49b19f6..000000000 --- a/scripts/pyinfra-complete.zsh +++ /dev/null @@ -1,28 +0,0 @@ -_pyinfra_completion() { - local -a completions - local -a completions_with_descriptions - local -a response - response=("${(@f)$( env COMP_WORDS="${words[*]}" \ - COMP_CWORD=$((CURRENT-1)) \ - _PYINFRA_COMPLETE="complete_zsh" \ - pyinfra )}") - - for key descr in ${(kv)response}; do - if [[ "$descr" == "_" ]]; then - completions+=("$key") - else - completions_with_descriptions+=("$key":"$descr") - fi - done - - if [ -n "$completions_with_descriptions" ]; then - _describe -V unsorted completions_with_descriptions -U -Q - fi - - if [ -n "$completions" ]; then - compadd -U -V unsorted -Q -a completions - fi - compstate[insert]="automenu" -} - -compdef _pyinfra_completion pyinfra; 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/progress.py b/src/pyinfra/progress.py index 1b5b37f17..0ce52507e 100644 --- a/src/pyinfra/progress.py +++ b/src/pyinfra/progress.py @@ -1,130 +1,69 @@ -import math 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 +from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn -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]) - - -def _print_spinner(stop_event, progress_queue): - if not IS_TTY or os.environ.get("PYINFRA_PROGRESS") == "off": - return - - progress = "" - text = "" - - while True: - # Stop when asked too - if stop_event.is_set(): - break - - WAIT_CHARS.rotate(1) - - try: - progress = progress_queue[-1] - except IndexError: - pass - - text = f" {' '.join((WAIT_CHARS[0], progress))}" - text = f"{text}\r" - - sys.stderr.write(text) - sys.stderr.flush() - - # 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) +from pyinfra.api.output import get_console, is_output_active @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 + """ + Display a Rich progress spinner while ``items`` are completed. + + Yields a ``progress(complete_item)`` callback; callers may ignore it (using + the spinner purely as a "busy" indicator). The display is refreshed + manually (``auto_refresh=False``) from the callback to stay well-behaved + under gevent (no background refresh greenlet). + """ + # If there's no active output we're not in CLI mode, so return a noop # handler and exit. if not is_output_active(): yield lambda complete_item: None return + # Allow disabling the spinner entirely. + if os.environ.get("PYINFRA_PROGRESS") == "off": + yield lambda complete_item: None + return + 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})", - ) + console = get_console() - if prefix_message: - message_bits.append(prefix_message) + columns = [ + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + ] + if total_items > 1: + columns.append(BarColumn()) + columns.append(TextColumn("{task.completed}/{task.total}")) - 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 + progress_bar = Progress( + *columns, + console=console, + transient=True, + auto_refresh=False, + ) - 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]}...}}" + description = prefix_message or "Working" + task_id = progress_bar.add_task(description, total=total_items) - message_bits.append(items_string) - - return " - ".join(message_bits) - - progress_queue = deque((make_progress_message(),)) + progress_bar.start() + progress_bar.refresh() def progress(complete_item): 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()) + progress_bar.update(task_id, advance=1) + progress_bar.refresh() - # 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 - - # Finally, stop the spinner - stop_event.set() - spinner_greenlet.join() + try: + yield progress + finally: + progress_bar.stop() diff --git a/src/pyinfra_cli/cli.py b/src/pyinfra_cli/cli.py index b324a77b0..f91efdc14 100644 --- a/src/pyinfra_cli/cli.py +++ b/src/pyinfra_cli/cli.py @@ -3,10 +3,12 @@ import warnings from fnmatch import fnmatch from getpass import getpass +from typing import Annotated from collections.abc import Iterable from os import chdir as os_chdir, environ, getcwd, path -import click +from cyclopts import App, Parameter +from rich.prompt import Confirm from pyinfra import __version__, logger, state from pyinfra.api import Config, Host, Inventory, State @@ -19,7 +21,10 @@ from pyinfra.context import ctx_config, ctx_inventory, ctx_state from pyinfra.operations import server +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 @@ -38,258 +43,278 @@ from .util import exec_file, load_deploy_file, load_func, parse_cli_arg from .virtualenv import init_virtualenv +# A repeatable flag: each occurrence appends a ``True``; the count is the verbosity level. +CountFlag = Annotated[list[bool], Parameter(negative="")] -def _exit() -> None: - if ctx_state.isset() and state.failed_hosts: - sys.exit(1) - sys.exit(0) +_FALSY_BOOL_VALUES = ("", "0", "n", "no", "off", "false", "f") +_TRUTHY_BOOL_VALUES = ("1", "y", "yes", "on", "true", "t") -def _print_support(ctx, param, value): - if not value: - return +def _lenient_bool(type_, tokens) -> bool: + """Lenient boolean conversion for environment variables. + + An empty value counts as unset (False) and ``on``/``off`` are accepted, + matching the legacy Click behaviour for ``PYINFRA_YES``. + """ + value = tokens[0].value.strip().lower() if tokens else "" + if value in _FALSY_BOOL_VALUES: + return False + if value in _TRUTHY_BOOL_VALUES: + return True + raise ValueError(f"invalid boolean value: {tokens[0].value!r}") + - logger.info("--> Support information:") +app = App( + name="pyinfra", + version=__version__, + version_flags=["--version"], + help_flags=["-h", "--help"], + console=stdout_console, + error_console=console, +) + +# Enable ``pyinfra --install-completion`` for shell autocompletion. +app.register_install_completion_command() + + +@app.command(name="--support") +def _support_command() -> None: + """Print useful information for support and exit.""" print_support_info() - ctx.exit() -CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"]) +def _exit() -> None: + if ctx_state.isset() and state.failed_hosts: + sys.exit(1) + sys.exit(0) -@click.command(context_settings=CONTEXT_SETTINGS) -@click.argument("inventory", nargs=1, type=click.Path(exists=False)) -@click.argument("operations", nargs=-1, required=True, type=click.Path(exists=False)) -@click.option( - "verbosity", - "-v", - count=True, - help="Print meta (-v), input (-vv) and output (-vvv).", -) -@click.option( - "--dry", - is_flag=True, - default=False, - help="Don't execute operations on the target hosts.", -) -@click.option( - "--diff", - is_flag=True, - default=False, - help="Show the differences when changing text files and templates.", -) -@click.option( - "-y", - "--yes", - is_flag=True, - default=False, - help="Execute operations immediately on hosts without prompt or checking for changes.", - envvar="PYINFRA_YES", - show_envvar=True, -) -@click.option( - "--limit", - help="Restrict the target hosts by name and group name.", - multiple=True, -) -@click.option( - "--exclude", - help="Exclude target hosts by name and group name.", - multiple=True, -) -@click.option("--fail-percent", type=int, help="% of hosts that need to fail before exiting early.") -@click.option( - "--data", - multiple=True, - help="Override data values, format key=value.", -) -@click.option( - "--group-data", - multiple=True, - help="Paths to load additional group data from (overrides matching keys).", -) -@click.option( - "--config", - "config_filename", - help="Specify config file to use (default: config.py).", - default="config.py", -) -@click.option( - "--chdir", - help="Set the working directory before executing.", -) -# Auth args -@click.option( - "--sudo", - is_flag=True, - default=False, - help="Whether to execute operations with sudo.", -) -@click.option("--sudo-user", help="Which user to sudo when sudoing.") -@click.option( - "--same-sudo-password", - is_flag=True, - default=False, - help="All hosts have the same sudo password, so ask only once.", -) -@click.option( - "--use-sudo-password", - is_flag=True, - default=False, - help="Whether to use a password with sudo.", -) -@click.option( - "--use-sudo-login", - is_flag=True, - default=False, - help="Use a login shell when sudo-ing.", -) -@click.option("--su-user", help="Which user to su to.") -@click.option( - "--dzdo", - is_flag=True, - default=False, - help="Whether to execute operations with dzdo.", -) -@click.option("--dzdo-user", help="Which user to dzdo when using dzdo.") -@click.option("--shell-executable", help='Shell to use (ex: "sh", "cmd", "ps").') -# Operation flow args -@click.option("--parallel", type=int, help="Number of operations to run in parallel.") -@click.option( - "--no-wait", - is_flag=True, - default=False, - help="Don't wait between operations for hosts.", -) -@click.option( - "--serial", - is_flag=True, - default=False, - help="Run operations in serial, host by host.", -) -@click.option( - "--retry", - type=int, - default=0, - help="Number of times to retry failed operations.", -) -@click.option( - "--retry-delay", - type=int, - default=5, - help="Delay in seconds between retry attempts.", -) -# SSH connector args -# TODO: remove the non-ssh-prefixed variants -@click.option("--ssh-user", "--user", "ssh_user", help="SSH user to connect as.") -@click.option("--ssh-port", "--port", "ssh_port", type=int, help="SSH port to connect to.") -@click.option("--ssh-key", "--key", "ssh_key", type=click.Path(), help="SSH Private key filename.") -@click.option( - "--ssh-key-password", - "--key-password", - "ssh_key_password", - help="SSH Private key password.", -) -@click.option("--ssh-password", "--password", "ssh_password", help="SSH password.") -@click.option( - "--ssh-password-prompt", - is_flag=True, - default=False, - help="Prompt for SSH password instead of passing it on the command line.", -) -# Eager commands (pyinfra --support) -@click.option( - "--support", - is_flag=True, - is_eager=True, - callback=_print_support, - help="Print useful information for support and exit.", -) -# Debug args -@click.option( - "--debug", - is_flag=True, - default=False, - help="Print debug logs from pyinfra.", -) -@click.option( - "--debug-all", - is_flag=True, - default=False, - help="Print debug logs from all packages including pyinfra.", -) -@click.option( - "--debug-facts", - is_flag=True, - default=False, - help="Print facts after generating operations and exit.", -) -@click.option( - "--debug-operations", - is_flag=True, - default=False, - help="Print operations after generating and exit.", -) -@click.option( - "--json", - "json_output", - is_flag=True, - default=False, - help=( - "Emit pure JSON output on stdout (for facts, debug-inventory, " - "debug-operations, dry runs and deploy results)." - ), -) -@click.version_option( - version=__version__, - prog_name="pyinfra", - message="%(prog)s: v%(version)s", -) -def cli(*args, **kwargs): - """ - pyinfra manages the state of one or more servers. It can be used for - app/service deployment, config management and ad-hoc command execution. +class CliCommands: + DEBUG_INVENTORY = "DEBUG_INVENTORY" + FACT = "FACT" + SHELL = "SHELL" + DEPLOY_FILES = "DEPLOY_FILES" + FUNC = "FUNC" + - Documentation: docs.pyinfra.com +@app.default +def cli( + inventory: str, + *operations: str, + verbose: Annotated[CountFlag, Parameter(name="-v")] = [], + dry: bool = False, + diff: bool = False, + yes: Annotated[ + bool, + Parameter(name=["-y", "--yes"], env_var="PYINFRA_YES", converter=_lenient_bool), + ] = False, + limit: tuple[str, ...] = (), + exclude: tuple[str, ...] = (), + fail_percent: int | None = None, + data: tuple[str, ...] = (), + group_data: tuple[str, ...] = (), + config_filename: Annotated[str, Parameter(name="--config")] = "config.py", + chdir: str | None = None, + sudo: bool = False, + sudo_user: str | None = None, + same_sudo_password: bool = False, + use_sudo_password: bool = False, + use_sudo_login: bool = False, + su_user: str | None = None, + dzdo: bool = False, + dzdo_user: str | None = None, + shell_executable: str | None = None, + parallel: int | None = None, + no_wait: bool = False, + serial: bool = False, + retry: int = 0, + retry_delay: int = 5, + ssh_user: Annotated[str | None, Parameter(name=["--ssh-user", "--user"])] = None, + ssh_port: Annotated[int | None, Parameter(name=["--ssh-port", "--port"])] = None, + ssh_key: Annotated[str | None, Parameter(name=["--ssh-key", "--key"])] = None, + ssh_key_password: Annotated[ + str | None, Parameter(name=["--ssh-key-password", "--key-password"]) + ] = None, + ssh_password: Annotated[str | None, Parameter(name=["--ssh-password", "--password"])] = None, + ssh_password_prompt: bool = False, + support: bool = False, + debug: bool = False, + debug_all: bool = False, + debug_facts: bool = False, + debug_operations: bool = False, + json_output: Annotated[bool, Parameter(name="--json")] = False, +): + """pyinfra manages the state of one or more servers. - # INVENTORY + It can be used for app/service deployment, config management and ad-hoc + command execution. Documentation: docs.pyinfra.com - \b - + a file (inventory.py) - + hostname (host.net) - + Comma separated hostnames: - host-1.net,host-2.net,@local + INVENTORY is a file (inventory.py), a hostname (host.net) or comma separated + hostnames (host-1.net,host-2.net,@local). - # OPERATIONS + Examples: - \b + ``` # Run one or more deploys against the inventory pyinfra INVENTORY deploy_web.py [deploy_db.py]... - \b # Run a single operation against the inventory pyinfra INVENTORY server.user pyinfra home=/home/pyinfra - \b # Execute an arbitrary command against the inventory pyinfra INVENTORY exec -- echo "hello world" - \b # Run one or more facts against the inventory pyinfra INVENTORY fact server.LinuxName [server.Users]... pyinfra INVENTORY fact files.File path=/path/to/file... - \b # Debug the inventory hosts and data pyinfra INVENTORY debug-inventory + ``` + + Parameters + ---------- + inventory + Inventory file, hostname(s) or connector to target. + operations + Deploy file(s), an operation + args, `exec -- command`, `fact ...` or + `debug-inventory`. + verbose + Print meta (-v), input (-vv) and output (-vvv). + dry + Don't execute operations on the target hosts. + diff + Show the differences when changing text files and templates. + yes + Execute operations immediately without prompt or checking for changes. + limit + Restrict the target hosts by name and group name. + exclude + Exclude target hosts by name and group name. + fail_percent + % of hosts that need to fail before exiting early. + data + Override data values, format key=value. + group_data + Paths to load additional group data from (overrides matching keys). + config_filename + Specify config file to use (default: config.py). + chdir + Set the working directory before executing. + sudo + Whether to execute operations with sudo. + sudo_user + Which user to sudo when sudoing. + same_sudo_password + All hosts have the same sudo password, so ask only once. + use_sudo_password + Whether to use a password with sudo. + use_sudo_login + Use a login shell when sudo-ing. + su_user + Which user to su to. + dzdo + Whether to execute operations with dzdo. + dzdo_user + Which user to dzdo when using dzdo. + shell_executable + Shell to use (ex: "sh", "cmd", "ps"). + parallel + Number of operations to run in parallel. + no_wait + Don't wait between operations for hosts. + serial + Run operations in serial, host by host. + retry + Number of times to retry failed operations. + retry_delay + Delay in seconds between retry attempts. + ssh_user + SSH user to connect as. + ssh_port + SSH port to connect to. + ssh_key + SSH Private key filename. + ssh_key_password + SSH Private key password. + ssh_password + SSH password. + ssh_password_prompt + Prompt for SSH password instead of passing it on the command line. + support + Print useful information for support and exit. + debug + Print debug logs from pyinfra. + debug_all + Print debug logs from all packages including pyinfra. + debug_facts + Print facts after generating operations and exit. + debug_operations + Print operations after generating and exit. + json_output + Emit pure JSON output on stdout (for facts, debug-inventory, + debug-operations, dry runs and deploy results). """ + if support: + logger.info("--> Support information:") + print_support_info() + return + + if not operations: + raise CliError( + "No operations provided.\n\n" + " Operation usage:\n" + " pyinfra INVENTORY deploy_web.py [deploy_db.py]...\n" + " pyinfra INVENTORY server.user pyinfra home=/home/pyinfra\n" + ' pyinfra INVENTORY exec -- echo "hello world"\n' + " pyinfra INVENTORY fact server.LinuxName [server.Users]..." + ) try: - _main(*args, **kwargs) + _main( + inventory=inventory, + operations=list(operations), + verbosity=len(verbose), + chdir=chdir, + ssh_user=ssh_user, + ssh_port=ssh_port, + ssh_key=ssh_key, + ssh_key_password=ssh_key_password, + ssh_password=ssh_password, + ssh_password_prompt=ssh_password_prompt, + same_sudo_password=same_sudo_password, + shell_executable=shell_executable, + sudo=sudo, + sudo_user=sudo_user, + use_sudo_password=use_sudo_password, + use_sudo_login=use_sudo_login, + su_user=su_user, + dzdo=dzdo, + dzdo_user=dzdo_user, + parallel=parallel, + fail_percent=fail_percent, + data=data, + group_data=group_data, + config_filename=config_filename, + dry=dry, + diff=diff, + yes=yes, + limit=limit, + exclude=exclude, + no_wait=no_wait, + serial=serial, + retry=retry, + retry_delay=retry_delay, + debug=debug, + debug_all=debug_all, + debug_facts=debug_facts, + debug_operations=debug_operations, + json_output=json_output, + ) except (CliError, UnexpectedExternalError): raise except PyinfraError as e: - # Re-raise "expected" pyinfra exceptions with our click exception wrapper + # Re-raise "expected" pyinfra exceptions with our exception wrapper raise WrappedError(e) except Exception as e: # Re-raise any unexpected internal exceptions as UnexpectedInternalError @@ -301,36 +326,28 @@ def cli(*args, **kwargs): disconnect_all(state) -class CliCommands: - DEBUG_INVENTORY = "DEBUG_INVENTORY" - FACT = "FACT" - SHELL = "SHELL" - DEPLOY_FILES = "DEPLOY_FILES" - FUNC = "FUNC" - - def _main( inventory, operations: list | tuple, verbosity: int, - chdir: str, - ssh_user, - ssh_port: int, - ssh_key, - ssh_key_password: str, - ssh_password: str, + chdir: str | None, + ssh_user: str | None, + ssh_port: int | None, + ssh_key: str | None, + ssh_key_password: str | None, + ssh_password: str | None, ssh_password_prompt: bool, same_sudo_password: bool, - shell_executable, + shell_executable: str | None, sudo: bool, - sudo_user: str, + sudo_user: str | None, use_sudo_password: bool, use_sudo_login: bool, - su_user: str, + su_user: str | None, dzdo: bool, - dzdo_user: str, - parallel: int, - fail_percent: int, + dzdo_user: str | None, + parallel: int | None, + fail_percent: int | None, data, group_data, config_filename: str, @@ -348,7 +365,6 @@ def _main( debug_facts: bool, debug_operations: bool, json_output: bool = False, - support: 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 @@ -458,13 +474,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 @@ -505,29 +520,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 @@ -882,7 +883,7 @@ def _prepare_deploy_operations(state, config, operations): for i, filename in enumerate(operations): config.lock_current_state() - _log_styled_msg = click.style(filename, bold=True) + _log_styled_msg = format_text(filename, bold=True) logger.info(f"Loading: {_log_styled_msg}") state.current_op_file_number = i 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 1f0e8e84e..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 cli +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: @@ -42,4 +40,9 @@ def _handle_interrupt(signum, frame): gevent.signal(signal.SIGINT, gevent.kill) signal.signal(signal.SIGINT, _handle_interrupt) # print the message and exit main - cli() + + try: + app() + except CliException as e: + e.show() + sys.exit(1) diff --git a/src/pyinfra_cli/prints.py b/src/pyinfra_cli/prints.py index 625017aef..b09fdf498 100644 --- a/src/pyinfra_cli/prints.py +++ b/src/pyinfra_cli/prints.py @@ -7,11 +7,11 @@ from typing import TYPE_CHECKING from collections.abc import Callable, Iterator -import click - from pyinfra import __version__, logger from pyinfra.api.host import Host +from pyinfra.api.output import format_text +from .console import console from .util import json_encode if TYPE_CHECKING: @@ -54,7 +54,8 @@ def jsonify(data, *args, **kwargs): def print_json(payload) -> None: - click.echo(jsonify(payload, default=json_encode)) + # Pure JSON on stdout — bypass Rich to guarantee byte-for-byte output. + print(jsonify(payload, default=json_encode)) def _host_to_dict(host: Host) -> dict: @@ -192,23 +193,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 +222,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 +231,29 @@ 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 print_inventory(state: State): 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) + console.print() + console.print(host.print_prefix) + console.print(f"--> Groups: {', '.join(host.groups)}") + console.print("--> Data:") + console.print(jsonify(host.data, indent=4, default=json_encode)) 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 +263,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 +271,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 +284,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 d717382d3..2cd01756b 100644 --- a/src/pyinfra_cli/util.py +++ b/src/pyinfra_cli/util.py @@ -11,10 +11,10 @@ from types import CodeType, FunctionType, ModuleType from collections.abc import Callable -import click import gevent from pyinfra import logger, state +from pyinfra.api.output import format_text from pyinfra.api.command import PyinfraCommand from pyinfra.api.exceptions import PyinfraError from pyinfra.api.host import HostData @@ -220,7 +220,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 dad845c97..75316375a 100644 --- a/src/pyinfra_cli/virtualenv.py +++ b/src/pyinfra_cli/virtualenv.py @@ -1,10 +1,10 @@ import os import sys -import click - from pyinfra import logger +from .console import console + def init_virtualenv() -> None: """ @@ -56,7 +56,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 = os.path.join( 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_exceptions.py b/tests/test_cli/test_cli_exceptions.py index b56fb6402..18dab5eb1 100644 --- a/tests/test_cli/test_cli_exceptions.py +++ b/tests/test_cli/test_cli_exceptions.py @@ -3,23 +3,17 @@ from unittest import TestCase import pytest -from click.testing import CliRunner from pyinfra.api import OperationError from pyinfra.api.exceptions import ArgumentTypeError from pyinfra_cli.exceptions import CliError, UnexpectedExternalError, WrappedError -from pyinfra_cli.main import cli from .util import run_cli class TestCliExceptions(TestCase): - @classmethod - def setUpClass(cls): - cls.runner = CliRunner() - def assert_cli_exception(self, args, message): - result = self.runner.invoke(cli, args, standalone_mode=False) + result = run_cli(*args) self.assertIsInstance(result.exception, CliError) assert getattr(result.exception, "message") == message diff --git a/tests/test_cli/util.py b/tests/test_cli/util.py index 8d3377d5c..94d9cf4b2 100644 --- a/tests/test_cli/util.py +++ b/tests/test_cli/util.py @@ -1,16 +1,65 @@ +import contextlib +from io import StringIO from os import chdir, getcwd -from click.testing import CliRunner - import pyinfra -from pyinfra_cli.main import cli +import pyinfra_cli.console as cli_console +from pyinfra_cli.cli import app +from pyinfra_cli.exceptions import CliException + + +class CliResult: + """Mimics the ``click.testing.Result`` interface used across the CLI tests.""" + + def __init__(self, exit_code, stdout, stderr, exception): + self.exit_code = exit_code + self.stdout = stdout + self.stderr = stderr + self.output = stdout + self.exception = exception def run_cli(*arguments): cwd = getcwd() pyinfra.is_cli = True - runner = CliRunner() - result = runner.invoke(cli, arguments, standalone_mode=False) - pyinfra.is_cli = False - chdir(cwd) - return result + + 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): + 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 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) + + return CliResult( + exit_code=exit_code, + stdout=stdout_buffer.getvalue(), + stderr=stderr_buffer.getvalue(), + exception=exception, + ) diff --git a/uv.lock b/uv.lock index 3edf34a9b..86ed6bb72 100644 --- a/uv.lock +++ b/uv.lock @@ -29,6 +29,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/25/8a/c46dcc25341b5bce5472c718902eb3d38600a903b14fa6aeecef3f21a46f/asttokens-3.0.0-py3-none-any.whl", hash = "sha256:e3078351a059199dd5138cb1c706e6430c05eff2ff136af5eb4790f9d28932e2", size = 26918, upload-time = "2024-11-30T04:30:10.946Z" }, ] +[[package]] +name = "attrs" +version = "26.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/8e/82a0fe20a541c03148528be8cac2408564a6c9a0cc7e9171802bc1d26985/attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32", size = 952055, upload-time = "2026-03-19T14:22:25.026Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, +] + [[package]] name = "bcrypt" version = "5.0.0" @@ -485,6 +494,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d9/b5/c5e179772ec38adb1c072b3aa13937d2860509ba32b2462bf1dda153833b/cryptography-46.0.2-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c4b93af7920cdf80f71650769464ccf1fb49a4b56ae0024173c24c48eb6b1612", size = 3438518, upload-time = "2025-10-01T00:29:06.139Z" }, ] +[[package]] +name = "cyclopts" +version = "4.20.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "docstring-parser" }, + { name = "rich" }, + { name = "rich-rst" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/74/89/f4c775c651d91f9cd8149f70baec94c902a34e5f17a7a67881881bcfb244/cyclopts-4.20.0.tar.gz", hash = "sha256:1d819de2b12dc6b1c9f17ce0f4937d82922c0b83ac846eb4b3289c9c9f321c9f", size = 190236, upload-time = "2026-06-29T15:04:42.253Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dd/d4/7243bc65d33c5ff5150569c42c8c1154aadae20b25638afe35f7f568489c/cyclopts-4.20.0-py3-none-any.whl", hash = "sha256:0b4337e9c11303d86b33d3f37c629dc01638f84591681e0e5611286bdd507646", size = 229383, upload-time = "2026-06-29T15:04:40.621Z" }, +] + [[package]] name = "decorator" version = "5.2.1" @@ -512,6 +538,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" }, ] +[[package]] +name = "docstring-parser" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/4d/f332313098c1de1b2d2ff91cf2674415cc7cddab2ca1b01ae29774bd5fdf/docstring_parser-0.18.0.tar.gz", hash = "sha256:292510982205c12b1248696f44959db3cdd1740237a968ea1e2e7a900eeb2015", size = 29341, upload-time = "2026-04-14T04:09:19.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, +] + +[[package]] +name = "docutils" +version = "0.23" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/39/a4/5180d9afc57e8fca05601dd652bdff19604c218814037fe90ffc7625a50a/docutils-0.23.tar.gz", hash = "sha256:746f5060322511280a1e50eb76846ed6bf2342984b2ac04dc42caa1a8d78799e", size = 2303823, upload-time = "2026-05-27T17:41:06.934Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/91/30151a39f7570f448ed84529390628a651d7f27c87d73c9b887f8189695e/docutils-0.23-py3-none-any.whl", hash = "sha256:25d013af9bf23bc1c7b2b093dff4208166c53a94786c9e447808335ef1185fea", size = 634701, upload-time = "2026-05-27T17:40:58.442Z" }, +] + [[package]] name = "exceptiongroup" version = "1.3.0" @@ -823,6 +867,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/1f/77fa3081e4f66ca3576c896ae5d31c3002ac6607f9747d2e3aa49227e464/markdown-3.10.2-py3-none-any.whl", hash = "sha256:e91464b71ae3ee7afd3017d9f358ef0baf158fd9a298db92f1d4761133824c36", size = 108180, upload-time = "2026-02-09T14:57:25.787Z" }, ] +[[package]] +name = "markdown-it-py" +version = "4.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, +] + [[package]] name = "markupsafe" version = "3.0.3" @@ -920,6 +976,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8f/8e/9ad090d3553c280a8060fbf6e24dc1c0c29704ee7d1c372f0c174aa59285/matplotlib_inline-0.1.7-py3-none-any.whl", hash = "sha256:df192d39a4ff8f21b1895d72e6a13f5fcc5099f00fa84384e0ea28c2cc0653ca", size = 9899, upload-time = "2024-04-15T13:44:43.265Z" }, ] +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + [[package]] name = "mergedeep" version = "1.3.4" @@ -1333,7 +1398,7 @@ wheels = [ name = "pyinfra" source = { editable = "." } dependencies = [ - { name = "click" }, + { name = "cyclopts" }, { name = "distro" }, { name = "gevent" }, { name = "jinja2" }, @@ -1341,6 +1406,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'" }, @@ -1348,7 +1414,6 @@ dependencies = [ [package.dev-dependencies] dev = [ - { name = "click" }, { name = "coverage" }, { name = "freezegun" }, { name = "ipdb" }, @@ -1376,7 +1441,6 @@ docs = [ { name = "zensical" }, ] test = [ - { name = "click" }, { name = "coverage" }, { name = "freezegun" }, { name = "mypy" }, @@ -1394,7 +1458,7 @@ 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" }, { name = "jinja2", specifier = ">3,<4" }, @@ -1402,6 +1466,7 @@ requires-dist = [ { name = "paramiko", specifier = ">=2.11,<5" }, { 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,<5" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, @@ -1409,7 +1474,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" }, @@ -1436,7 +1500,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" }, @@ -1662,6 +1725,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, ] +[[package]] +name = "rich" +version = "15.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, +] + +[[package]] +name = "rich-rst" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "docutils" }, + { name = "rich" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bc/6d/a506aaa4a9eaa945ed8ab2b7347859f53593864289853c5d6d62b77246e0/rich_rst-1.3.2.tar.gz", hash = "sha256:a1196fdddf1e364b02ec68a05e8ff8f6914fee10fbca2e6b6735f166bb0da8d4", size = 14936, upload-time = "2025-10-14T16:49:45.332Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/13/2f/b4530fbf948867702d0a3f27de4a6aab1d156f406d72852ab902c4d04de9/rich_rst-1.3.2-py3-none-any.whl", hash = "sha256:a99b4907cbe118cf9d18b0b44de272efa61f15117c61e39ebdc431baf5df722a", size = 12567, upload-time = "2025-10-14T16:49:42.953Z" }, +] + [[package]] name = "ruff" version = "0.14.0" From abf48ff9e1f898fd2f7de16609ab10cdc078ffd1 Mon Sep 17 00:00:00 2001 From: "Axel H." Date: Sat, 4 Jul 2026 20:39:40 +0200 Subject: [PATCH 2/3] feat(cli): rich debug-inventory, help and argument improvements - debug-inventory renders a Rich table with syntax-highlighted, pretty-printed host data; --json output is highlighted on a terminal and plain when piped - de-duplicate host groups and tidy the table layout - make INVENTORY positional-only and organise options into ordered help groups - syntax-highlight the CLI help examples --- src/pyinfra_cli/cli.py | 191 ++++++++++++++++++++---------- src/pyinfra_cli/prints.py | 127 ++++++++++++++++++-- tests/test_cli/test_cli_prints.py | 87 ++++++++++++++ 3 files changed, 337 insertions(+), 68 deletions(-) create mode 100644 tests/test_cli/test_cli_prints.py diff --git a/src/pyinfra_cli/cli.py b/src/pyinfra_cli/cli.py index f91efdc14..95dae9b98 100644 --- a/src/pyinfra_cli/cli.py +++ b/src/pyinfra_cli/cli.py @@ -7,7 +7,7 @@ from collections.abc import Iterable from os import chdir as os_chdir, environ, getcwd, path -from cyclopts import App, Parameter +from cyclopts import App, Group, Parameter from rich.prompt import Confirm from pyinfra import __version__, logger, state @@ -64,18 +64,86 @@ 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=__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. app.register_install_completion_command() +# Parameter groups for the help page (ordered top-to-bottom as declared). +# ``negative=""`` disables the auto-generated ``--no-*`` / ``--empty-*`` flags to +# match the original flag-only CLI UX. +_no_negative = Parameter(negative="") +GROUP_EXECUTION = Group.create_ordered("Execution", default_parameter=_no_negative) +GROUP_INVENTORY = Group.create_ordered("Inventory & Data", default_parameter=_no_negative) +GROUP_PRIVILEGE = Group.create_ordered("Privilege Escalation", default_parameter=_no_negative) +GROUP_SSH = Group.create_ordered("SSH Connection", default_parameter=_no_negative) +GROUP_DEBUG = Group.create_ordered("Debugging & Output", default_parameter=_no_negative) + @app.command(name="--support") def _support_command() -> None: @@ -100,77 +168,80 @@ class CliCommands: @app.default def cli( inventory: str, + /, *operations: str, - verbose: Annotated[CountFlag, Parameter(name="-v")] = [], - dry: bool = False, - diff: bool = False, + # Execution + dry: Annotated[bool, Parameter(group=GROUP_EXECUTION)] = False, yes: Annotated[ bool, - Parameter(name=["-y", "--yes"], env_var="PYINFRA_YES", converter=_lenient_bool), + Parameter( + name=["-y", "--yes"], + env_var="PYINFRA_YES", + converter=_lenient_bool, + group=GROUP_EXECUTION, + ), ] = False, - limit: tuple[str, ...] = (), - exclude: tuple[str, ...] = (), - fail_percent: int | None = None, - data: tuple[str, ...] = (), - group_data: tuple[str, ...] = (), - config_filename: Annotated[str, Parameter(name="--config")] = "config.py", - chdir: str | None = None, - sudo: bool = False, - sudo_user: str | None = None, - same_sudo_password: bool = False, - use_sudo_password: bool = False, - use_sudo_login: bool = False, - su_user: str | None = None, - dzdo: bool = False, - dzdo_user: str | None = None, - shell_executable: str | None = None, - parallel: int | None = None, - no_wait: bool = False, - serial: bool = False, - retry: int = 0, - retry_delay: int = 5, - ssh_user: Annotated[str | None, Parameter(name=["--ssh-user", "--user"])] = None, - ssh_port: Annotated[int | None, Parameter(name=["--ssh-port", "--port"])] = None, - ssh_key: Annotated[str | None, Parameter(name=["--ssh-key", "--key"])] = None, + parallel: Annotated[int | None, Parameter(group=GROUP_EXECUTION)] = None, + no_wait: Annotated[bool, Parameter(group=GROUP_EXECUTION)] = False, + serial: Annotated[bool, Parameter(group=GROUP_EXECUTION)] = False, + fail_percent: Annotated[int | None, Parameter(group=GROUP_EXECUTION)] = None, + retry: Annotated[int, Parameter(group=GROUP_EXECUTION)] = 0, + retry_delay: Annotated[int, Parameter(group=GROUP_EXECUTION)] = 5, + shell_executable: Annotated[str | None, Parameter(group=GROUP_EXECUTION)] = None, + # Inventory & Data + limit: Annotated[tuple[str, ...], Parameter(group=GROUP_INVENTORY)] = (), + exclude: Annotated[tuple[str, ...], Parameter(group=GROUP_INVENTORY)] = (), + data: Annotated[tuple[str, ...], Parameter(group=GROUP_INVENTORY)] = (), + group_data: Annotated[tuple[str, ...], Parameter(group=GROUP_INVENTORY)] = (), + config_filename: Annotated[ + str, Parameter(name="--config", group=GROUP_INVENTORY) + ] = "config.py", + chdir: Annotated[str | None, Parameter(group=GROUP_INVENTORY)] = None, + # Privilege escalation + sudo: Annotated[bool, Parameter(group=GROUP_PRIVILEGE)] = False, + sudo_user: Annotated[str | None, Parameter(group=GROUP_PRIVILEGE)] = None, + same_sudo_password: Annotated[bool, Parameter(group=GROUP_PRIVILEGE)] = False, + use_sudo_password: Annotated[bool, Parameter(group=GROUP_PRIVILEGE)] = False, + use_sudo_login: Annotated[bool, Parameter(group=GROUP_PRIVILEGE)] = False, + su_user: Annotated[str | None, Parameter(group=GROUP_PRIVILEGE)] = None, + dzdo: Annotated[bool, Parameter(group=GROUP_PRIVILEGE)] = False, + dzdo_user: Annotated[str | None, Parameter(group=GROUP_PRIVILEGE)] = None, + # SSH connection + ssh_user: Annotated[ + str | None, Parameter(name=["--ssh-user", "--user"], group=GROUP_SSH) + ] = None, + ssh_port: Annotated[ + int | None, Parameter(name=["--ssh-port", "--port"], group=GROUP_SSH) + ] = None, + ssh_key: Annotated[str | None, Parameter(name=["--ssh-key", "--key"], group=GROUP_SSH)] = None, ssh_key_password: Annotated[ - str | None, Parameter(name=["--ssh-key-password", "--key-password"]) + str | None, Parameter(name=["--ssh-key-password", "--key-password"], group=GROUP_SSH) + ] = None, + ssh_password: Annotated[ + str | None, Parameter(name=["--ssh-password", "--password"], group=GROUP_SSH) ] = None, - ssh_password: Annotated[str | None, Parameter(name=["--ssh-password", "--password"])] = None, - ssh_password_prompt: bool = False, - support: bool = False, - debug: bool = False, - debug_all: bool = False, - debug_facts: bool = False, - debug_operations: bool = False, - json_output: Annotated[bool, Parameter(name="--json")] = False, + ssh_password_prompt: Annotated[bool, Parameter(group=GROUP_SSH)] = False, + # Debugging & output + verbose: Annotated[CountFlag, Parameter(name="-v", group=GROUP_DEBUG)] = [], + diff: Annotated[bool, Parameter(group=GROUP_DEBUG)] = False, + json_output: Annotated[bool, Parameter(name="--json", group=GROUP_DEBUG)] = False, + support: Annotated[bool, Parameter(group=GROUP_DEBUG)] = False, + debug: Annotated[bool, Parameter(group=GROUP_DEBUG)] = False, + debug_all: Annotated[bool, Parameter(group=GROUP_DEBUG)] = False, + debug_facts: Annotated[bool, Parameter(group=GROUP_DEBUG)] = False, + debug_operations: Annotated[bool, Parameter(group=GROUP_DEBUG)] = False, ): """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). - - 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" + command execution. - # Run one or more facts against the inventory - pyinfra INVENTORY fact server.LinuxName [server.Users]... - pyinfra INVENTORY fact files.File path=/path/to/file... + Documentation: [cyan][link=https://docs.pyinfra.com]docs.pyinfra.com[/link][/cyan] - # 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 ---------- diff --git a/src/pyinfra_cli/prints.py b/src/pyinfra_cli/prints.py index b09fdf498..4f3472e9f 100644 --- a/src/pyinfra_cli/prints.py +++ b/src/pyinfra_cli/prints.py @@ -7,11 +7,16 @@ from typing import TYPE_CHECKING from collections.abc import Callable, Iterator +from rich.console import Group +from rich.json import JSON +from rich.padding import Padding +from rich.table import Table + from pyinfra import __version__, logger from pyinfra.api.host import Host from pyinfra.api.output import format_text -from .console import console +from .console import console, stdout_console from .util import json_encode if TYPE_CHECKING: @@ -53,9 +58,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: - # Pure JSON on stdout — bypass Rich to guarantee byte-for-byte output. - print(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: @@ -240,13 +286,78 @@ def print_fact(fact_data): 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: - console.print() - console.print(host.print_prefix) - console.print(f"--> Groups: {', '.join(host.groups)}") - console.print("--> Data:") - console.print(jsonify(host.data, indent=4, default=json_encode)) + # 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): 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")) == "" From c6eb31b831289ac86a1877470b90a49ae54b1aeb Mon Sep 17 00:00:00 2001 From: "Axel H." Date: Sat, 4 Jul 2026 21:06:19 +0200 Subject: [PATCH 3/3] feat(cli): live hierarchical progress tree and informative error handling Replace the flat `--> ` phase logs with a Rich rendering of deploys. - Drop the `--> ` log prefix; render proposed changes and results as Rich trees/tables (file -> operation -> hosts) with per-op success/error/no-change counts and failed hosts listed under errored operations - Hierarchical live tree (TTY, non-JSON, all verbosity levels): phases (Loading, Connecting, Preparing, each operation) render as spinner -> check/cross rows with a parent progress bar (n/total); verbose detail lines (facts, command input/output) nest under the host nodes - New pyinfra_cli.routing attributes log/echo lines to hosts and records per-host errors so nothing interleaves with the live region - Deferred, aggregated failure prompt: during Prepare, hosts evaluate in parallel, so failures are collected and a single prompt (preceded by a red "Failed hosts" block with each host's error) is shown after the phase - Consistent status colours; dim the `@connector/` host-name prefix; restyle the bracket-free host log prefix - New api operation_host_skipped state callback; fire operation_end before fail_hosts so the tree finalises before prompting In piped/JSON modes the live tree is disabled and output stays clean/pure. --- src/pyinfra/api/connect.py | 4 +- src/pyinfra/api/facts.py | 7 +- src/pyinfra/api/host.py | 17 +- src/pyinfra/api/operations.py | 19 +- src/pyinfra/api/state.py | 4 + src/pyinfra/api/util.py | 2 +- src/pyinfra/progress.py | 111 ++++-- src/pyinfra_cli/cli.py | 217 +++++++---- src/pyinfra_cli/console.py | 54 ++- src/pyinfra_cli/exceptions.py | 12 +- src/pyinfra_cli/log.py | 52 ++- src/pyinfra_cli/prints.py | 208 +++++------ src/pyinfra_cli/progress.py | 571 +++++++++++++++++++++++++++++ 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 +- 18 files changed, 1222 insertions(+), 285 deletions(-) create mode 100644 src/pyinfra_cli/progress.py create mode 100644 src/pyinfra_cli/routing.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 9d7dc62ba..1daabd91e 100644 --- a/src/pyinfra/api/facts.py +++ b/src/pyinfra/api/facts.py @@ -177,7 +177,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..95781d8f0 100644 --- a/src/pyinfra/api/host.py +++ b/src/pyinfra/api/host.py @@ -198,15 +198,26 @@ 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}") 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/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 a48b5f58e..da19822e3 100644 --- a/src/pyinfra/api/util.py +++ b/src/pyinfra/api/util.py @@ -195,7 +195,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/progress.py b/src/pyinfra/progress.py index 0ce52507e..277b3db3a 100644 --- a/src/pyinfra/progress.py +++ b/src/pyinfra/progress.py @@ -1,60 +1,96 @@ +from __future__ import annotations + import os from contextlib import contextmanager +from typing import TYPE_CHECKING, Any -from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn +from rich.errors import LiveError +from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn from pyinfra.api.output import get_console, is_output_active +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Iterator + +# 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 + + +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 + + +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" + + +def _noop_progress(complete_item: Any) -> None: + pass + @contextmanager -def progress_spinner(items, prefix_message=None): +def progress_spinner( + items: Iterable[Any], + prefix_message: str | None = None, +) -> Iterator[Callable[[Any], None]]: """ - Display a Rich progress spinner while ``items`` are completed. + Display a Rich progress bar while ``items`` are completed. Yields a ``progress(complete_item)`` callback; callers may ignore it (using - the spinner purely as a "busy" indicator). The display is refreshed + 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 there's no active output we're not in CLI mode, so return a noop - # handler and exit. - if not is_output_active(): - yield lambda complete_item: None + if not _spinner_enabled(): + yield _noop_progress return - # Allow disabling the spinner entirely. - if os.environ.get("PYINFRA_PROGRESS") == "off": - yield lambda complete_item: None - return + global _active_spinners, _progress if not isinstance(items, set): items = set(items) total_items = len(items) - console = get_console() - - columns = [ - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - ] - if total_items > 1: - columns.append(BarColumn()) - columns.append(TextColumn("{task.completed}/{task.total}")) - - progress_bar = Progress( - *columns, - console=console, - transient=True, - auto_refresh=False, - ) + progress_bar = _get_progress() + + 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.start() progress_bar.refresh() - def progress(complete_item): + def progress(complete_item: Any) -> None: if complete_item not in items: raise ValueError( f"Invalid complete item: {complete_item} not in {items}", @@ -66,4 +102,15 @@ def progress(complete_item): try: yield progress finally: - progress_bar.stop() + # 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 95dae9b98..b3f0792f5 100644 --- a/src/pyinfra_cli/cli.py +++ b/src/pyinfra_cli/cli.py @@ -24,7 +24,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 @@ -134,6 +136,13 @@ def _build_usage() -> str: # Enable ``pyinfra --install-completion`` for shell autocompletion. app.register_install_completion_command() + +@app.command(name="--support") +def _support_command() -> None: + """Print useful information for support and exit.""" + print_support_info() + + # Parameter groups for the help page (ordered top-to-bottom as declared). # ``negative=""`` disables the auto-generated ``--no-*`` / ``--empty-*`` flags to # match the original flag-only CLI UX. @@ -145,12 +154,6 @@ def _build_usage() -> str: GROUP_DEBUG = Group.create_ordered("Debugging & Output", default_parameter=_no_negative) -@app.command(name="--support") -def _support_command() -> None: - """Print useful information for support and exit.""" - print_support_info() - - def _exit() -> None: if ctx_state.isset() and state.failed_hosts: sys.exit(1) @@ -327,7 +330,7 @@ def cli( debug-operations, dry runs and deploy results). """ if support: - logger.info("--> Support information:") + logger.info("Support information:") print_support_info() return @@ -391,8 +394,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) @@ -437,12 +443,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: @@ -463,27 +463,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: ") @@ -503,21 +521,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: @@ -528,22 +554,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( """ @@ -577,11 +626,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) @@ -602,6 +657,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): @@ -730,7 +814,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: @@ -822,9 +906,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 @@ -878,7 +969,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) @@ -889,16 +980,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..4510291d8 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,48 @@ from __future__ import annotations +from functools import cache from typing import Any from rich.console import Console 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, +) + + +@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 +59,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 +71,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 +82,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..d43571ef1 100644 --- a/src/pyinfra_cli/log.py +++ b/src/pyinfra_cli/log.py @@ -6,6 +6,7 @@ from pyinfra import logger, state from pyinfra.context import ctx_state +from . import routing from .console import console, format_text @@ -13,9 +14,40 @@ 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)) + # 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 (Connected/Ready/Loaded fact/...) only + # in verbose mode — at default verbosity the node status + # already conveys them. + if record.levelno >= logging.WARNING or tree.verbose: + 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 +80,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 +96,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/prints.py b/src/pyinfra_cli/prints.py index 4f3472e9f..9bb270b96 100644 --- a/src/pyinfra_cli/prints.py +++ b/src/pyinfra_cli/prints.py @@ -2,20 +2,22 @@ import json import platform -import re import sys -from typing import TYPE_CHECKING -from collections.abc import Callable, Iterator +from typing import TYPE_CHECKING, Any +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 @@ -23,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]] = {} @@ -240,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] @@ -364,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) @@ -408,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] @@ -465,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 = [] @@ -482,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 magenta")) + 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} @@ -531,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 magenta")) + 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..92eebfd6c --- /dev/null +++ b/src/pyinfra_cli/progress.py @@ -0,0 +1,571 @@ +""" +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 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.state import BaseStateCallback + +from . import routing +from .console import console + +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 magenta. +_BAR_STYLE = { + NodeKind.FILE: "magenta", + 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 + + +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 + # Verbose per-host lines (fact loads, command input/output, ...) + # rendered nested under this node: list of (text, is_error). + self.detail_lines: list[tuple[str, bool]] = [] + 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((text, is_error)) + + def succeed(self, detail: str | None = None) -> None: + self.status = NodeStatus.OK + self.detail = detail + + 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 magenta") + + 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=_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. + if self.detail_lines: + lines = self.detail_lines + if self.status == NodeStatus.RUNNING: + lines = lines[-RUNNING_DETAIL_LINES:] + detail_indent = " " * (self.depth + 1) + for text, is_error in lines: + table.add_row( + Text(""), + Text(f"{detail_indent}{text}", style="red" if 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._live.stop() + + def _rotate_region(self) -> None: + """Persist the current region to scrollback and start a fresh one.""" + if self._live.is_started: + self._live.stop() + self._roots = [] + self._file_nodes = {} + self._live = self._new_live() + self._live.start() + + 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 + node.add_detail_line(text, is_error=is_error) + + # 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() + + 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 + + @override + def operation_host_start(self, state: State, host: Host, op_hash) -> None: # type: ignore[override] + parent = self._op_nodes.get(op_hash) + if parent is None: + return + 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: + node.succeed() + + @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/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 2cd01756b..0a95a08a1 100644 --- a/src/pyinfra_cli/util.py +++ b/src/pyinfra_cli/util.py @@ -214,6 +214,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()): @@ -225,18 +233,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"], )