pyinfra 3.5.1__py3-none-any.whl → 3.6.1__py3-none-any.whl

pyinfra/api/__init__.py

@@ -14,6 +14,7 @@ from .exceptions import (  # noqa: F401
     FactError,
     FactTypeError,
     FactValueError,
+    FactProcessError,
     InventoryError,
     OperationError,
     OperationTypeError,

pyinfra/api/arguments.py

@@ -70,12 +70,15 @@ class ConnectorArguments(TypedDict, total=False):
     _success_exit_codes: Iterable[int]
     _timeout: int
     _get_pty: bool
-    _stdin: Union[str, Iterable[str]]
+    _stdin: Union[str, list[str], Iterable[str]]
 
     # Retry arguments
     _retries: int
     _retry_delay: Union[int, float]
-    _retry_until: Optional[Callable[[dict], bool]]
+    _retry_until: Callable[[dict], bool]
+
+    # Temp directory argument
+    _temp_dir: str
 
 
 def generate_env(config: "Config", value: dict) -> dict:
@@ -163,6 +166,10 @@ shell_argument_meta: dict[str, ArgumentMeta] = {
         "String or buffer to send to the stdin of any commands.",
         default=lambda _: None,
     ),
+    "_temp_dir": ArgumentMeta(
+        "Temporary directory on the remote host for file operations.",
+        default=lambda config: config.TEMP_DIR,
+    ),
 }
 
 

pyinfra/api/arguments_typed.py

@@ -1,16 +1,6 @@
 from __future__ import annotations
 
-from typing import (
-    TYPE_CHECKING,
-    Callable,
-    Generator,
-    Generic,
-    Iterable,
-    List,
-    Mapping,
-    Optional,
-    Union,
-)
+from typing import TYPE_CHECKING, Callable, Generator, Generic, Iterable, List, Mapping, Union
 
 from typing_extensions import ParamSpec, Protocol
 
@@ -36,36 +26,41 @@ class PyinfraOperation(Generic[P], Protocol):
         #
         # Auth
         _sudo: bool = False,
-        _sudo_user: Optional[str] = None,
+        _sudo_user: None | str = None,
         _use_sudo_login: bool = False,
-        _sudo_password: Optional[str] = None,
+        _sudo_password: None | str = None,
         _preserve_sudo_env: bool = False,
-        _su_user: Optional[str] = None,
+        _su_user: None | str = None,
         _use_su_login: bool = False,
         _preserve_su_env: bool = False,
-        _su_shell: Optional[str] = None,
+        _su_shell: None | str = None,
         _doas: bool = False,
-        _doas_user: Optional[str] = None,
+        _doas_user: None | str = None,
         # Shell arguments
-        _shell_executable: Optional[str] = None,
-        _chdir: Optional[str] = None,
-        _env: Optional[Mapping[str, str]] = None,
+        _shell_executable: None | str = None,
+        _chdir: None | str = None,
+        _env: None | Mapping[str, str] = None,
         # Connector control
         _success_exit_codes: Iterable[int] = (0,),
-        _timeout: Optional[int] = None,
+        _timeout: None | int = None,
         _get_pty: bool = False,
-        _stdin: Union[None, str, list[str], tuple[str, ...]] = None,
+        _stdin: None | Union[str, list[str], Iterable[str]] = None,
+        # Retry arguments
+        _retries: None | int = None,
+        _retry_delay: None | Union[int, float] = None,
+        _retry_until: None | Callable[[dict], bool] = None,
+        _temp_dir: None | str = None,
         #
         # MetaArguments
         #
-        name: Optional[str] = None,
+        name: None | str = None,
         _ignore_errors: bool = False,
         _continue_on_error: bool = False,
         _if: Union[List[Callable[[], bool]], Callable[[], bool], None] = None,
         #
         # ExecutionArguments
         #
-        _parallel: Optional[int] = None,
+        _parallel: None | int = None,
         _run_once: bool = False,
         _serial: bool = False,
         #

pyinfra/api/command.py

@@ -242,13 +242,19 @@ class FunctionCommand(PyinfraCommand):
             self.function(*self.args, **self.kwargs)
             return
 
-        def execute_function() -> None:
+        def execute_function() -> None | Exception:
             with ctx_config.use(state.config.copy()):
                 with ctx_host.use(host):
-                    self.function(*self.args, **self.kwargs)
+                    try:
+                        self.function(*self.args, **self.kwargs)
+                    except Exception as e:
+                        return e
+            return None
 
         greenlet = gevent.spawn(execute_function)
-        return greenlet.get()
+        exception = greenlet.get()
+        if exception is not None:
+            raise exception
 
 
 class RsyncCommand(PyinfraCommand):

pyinfra/api/deploy.py

@@ -41,7 +41,7 @@ def add_deploy(state: "State", deploy_func: Callable[..., Any], *args, **kwargs)
             ).format(get_call_location()),
         )
 
-    hosts = kwargs.pop("host", state.inventory.iter_active_hosts())
+    hosts = kwargs.pop("host", state.inventory.get_active_hosts())
     if isinstance(hosts, Host):
         hosts = [hosts]
 

pyinfra/api/exceptions.py

@@ -29,6 +29,12 @@ class FactValueError(FactError, ValueError):
     """
 
 
+class FactProcessError(FactError, RuntimeError):
+    """
+    Exception raised when the data gathered for a fact cannot be processed.
+    """
+
+
 class OperationError(PyinfraError):
     """
     Exception raised during fact gathering staging if an operation is unable to
@@ -48,6 +54,12 @@ class OperationValueError(OperationError, ValueError):
     """
 
 
+class NestedOperationError(OperationError):
+    """
+    Exception raised when a nested (immediately executed) operation fails.
+    """
+
+
 class DeployError(PyinfraError):
     """
     User exception for raising in deploys or sub deploys.

pyinfra/api/facts.py

@@ -14,7 +14,7 @@ import inspect
 import re
 from inspect import getcallargs
 from socket import error as socket_error, timeout as timeout_error
-from typing import TYPE_CHECKING, Any, Callable, Generic, Iterable, Optional, Type, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Type, TypeVar, cast
 
 import click
 import gevent
@@ -24,6 +24,7 @@ from typing_extensions import override
 from pyinfra import logger
 from pyinfra.api import StringCommand
 from pyinfra.api.arguments import all_global_arguments, pop_global_arguments
+from pyinfra.api.exceptions import FactProcessError
 from pyinfra.api.util import (
     get_kwargs_str,
     log_error_or_warning,
@@ -86,7 +87,7 @@ class FactBase(Generic[T]):
 
         return cast(T, None)
 
-    def process(self, output: Iterable[str]) -> T:
+    def process(self, output: list[str]) -> T:
         # NOTE: TypeVar does not support a default, so we have to cast this str -> T
         return cast(T, "\n".join(output))
 
@@ -151,7 +152,7 @@ def get_facts(state, *args, **kwargs):
     with ctx_state.use(state):
         greenlet_to_host = {
             state.pool.spawn(get_host_fact, host, *args, **kwargs): host
-            for host in state.inventory.iter_active_hosts()
+            for host in state.inventory.get_active_hosts()
         }
 
     results = {}
@@ -269,7 +270,22 @@ def _get_fact(
 
     if status:
         if stdout_lines:
-            data = fact.process(stdout_lines)
+            try:
+                data = fact.process(stdout_lines)
+            except FactProcessError as e:
+                log_error_or_warning(
+                    host,
+                    global_kwargs["_ignore_errors"],
+                    description=("could not process fact: {0} {1}").format(
+                        name, get_kwargs_str(fact_kwargs)
+                    ),
+                    exception=e,
+                )
+
+                # Check we've not failed
+                if apply_failed_hosts and not global_kwargs["_ignore_errors"]:
+                    state.fail_hosts({host})
+
     elif stderr_lines:
         # If we have error output and that error is sudo or su stating the user
         # does not exist, do not fail but instead return the default fact value.

pyinfra/api/host.py

@@ -328,6 +328,9 @@ class Host:
 
         return temp_directory
 
+    def get_temp_dir_config(self):
+        return self.state.config.TEMP_DIR or self.state.config.DEFAULT_TEMP_DIR
+
     def get_temp_filename(
         self,
         hash_key: Optional[str] = None,

pyinfra/api/inventory.py

@@ -158,11 +158,11 @@ class Inventory:
 
         return iter(self.hosts.values())
 
-    def iter_active_hosts(self) -> Iterator["Host"]:
+    def get_active_hosts(self) -> list["Host"]:
         """
         Iterates over active inventory hosts.
         """
-        return iter(self.state.active_hosts)
+        return list(self.state.active_hosts)
 
     def len_active_hosts(self) -> int:
         """

pyinfra/api/metadata.py

@@ -0,0 +1,69 @@
+"""
+Support parsing pyinfra-metadata.toml
+
+Currently just parses plugins and their metadata.
+"""
+
+import tomllib
+from typing import Literal, get_args
+
+from pydantic import BaseModel, TypeAdapter, field_validator
+
+AllowedTagType = Literal[
+    "boot",
+    "containers",
+    "database",
+    "service-management",
+    "package-manager",
+    "python",
+    "ruby",
+    "javascript",
+    "configuration-management",
+    "security",
+    "storage",
+    "system",
+    "rust",
+    "version-control-system",
+]
+
+
+class Tag(BaseModel):
+    """Representation of a plugin tag."""
+
+    value: AllowedTagType
+
+    @field_validator("value", mode="before")
+    def _validate_value(cls, v) -> AllowedTagType:
+        allowed_tags = set(get_args(AllowedTagType))
+        if v not in allowed_tags:
+            raise ValueError(f"Invalid tag: {v}. Allowed: {allowed_tags}")
+        return v
+
+    @property
+    def title_case(self) -> str:
+        return " ".join([t.title() for t in self.value.split("-")])
+
+
+ALLOWED_TAGS = [Tag(value=tag) for tag in set(get_args(AllowedTagType))]
+
+
+class Plugin(BaseModel):
+    """Representation of a pyinfra plugin."""
+
+    name: str
+    # description: str # FUTURE we should grab these from doc strings
+    path: str
+    type: Literal["operation", "fact", "connector", "deploy"]
+    tags: list[Tag]
+
+    @field_validator("tags", mode="before")
+    def _wrap_tags(cls, v):
+        return [Tag(value=tag) if not isinstance(tag, Tag) else tag for tag in v]
+
+
+def parse_plugins(metadata_text: str) -> list[Plugin]:
+    """Given the contents of a pyinfra-metadata.toml parse out the plugins."""
+    pyinfra_metadata = tomllib.loads(metadata_text).get("pyinfra", None)
+    if not pyinfra_metadata:
+        raise ValueError("Missing [pyinfra.plugins] section in pyinfra-metadata.toml")
+    return TypeAdapter(list[Plugin]).validate_python(pyinfra_metadata["plugins"].values())

pyinfra/api/operation.py

@@ -22,7 +22,7 @@ from pyinfra.context import ctx_host, ctx_state
 from .arguments import EXECUTION_KWARG_KEYS, AllArguments, pop_global_arguments
 from .arguments_typed import PyinfraOperation
 from .command import PyinfraCommand, StringCommand
-from .exceptions import OperationValueError, PyinfraError
+from .exceptions import NestedOperationError, OperationValueError, PyinfraError
 from .host import Host
 from .operations import run_host_op
 from .state import State, StateOperationHostData, StateOperationMeta, StateStage
@@ -221,7 +221,7 @@ def add_op(state: State, op_func, *args, **kwargs):
             ),
         )
 
-    hosts = kwargs.pop("host", state.inventory.iter_active_hosts())
+    hosts = kwargs.pop("host", state.inventory.get_active_hosts())
     if isinstance(hosts, Host):
         hosts = [hosts]
 
@@ -266,7 +266,9 @@ def _wrap_operation(func: Callable[P, Generator], _set_in_op: bool = True) -> Py
         state = context.state
         host = context.host
 
-        if state.current_stage < StateStage.Prepare or state.current_stage > StateStage.Execute:
+        if pyinfra.is_cli and (
+            state.current_stage < StateStage.Prepare or state.current_stage > StateStage.Execute
+        ):
             raise Exception("Cannot call operations outside of Prepare/Execute stages")
 
         if host.in_op:
@@ -470,8 +472,11 @@ def execute_immediately(state, host, op_hash):
     op_meta = state.get_op_meta(op_hash)
     op_data = state.get_op_data_for_host(host, op_hash)
     op_data.parent_op_hash = host.executing_op_hash
+
     log_operation_start(op_meta, op_types=["nested"], prefix="")
-    run_host_op(state, host, op_hash)
+
+    if run_host_op(state, host, op_hash) is False:
+        raise NestedOperationError(op_hash)
 
 
 def _get_arg_value(arg):

pyinfra/api/operations.py

@@ -4,7 +4,7 @@ import time
 import traceback
 from itertools import product
 from socket import error as socket_error, timeout as timeout_error
-from typing import TYPE_CHECKING, Optional, cast
+from typing import TYPE_CHECKING, cast
 
 import click
 import gevent
@@ -17,7 +17,7 @@ from pyinfra.progress import progress_spinner
 
 from .arguments import CONNECTOR_ARGUMENT_KEYS, ConnectorArguments
 from .command import FunctionCommand, PyinfraCommand, StringCommand
-from .exceptions import PyinfraError
+from .exceptions import NestedOperationError, PyinfraError
 from .util import (
     format_exception,
     log_error_or_warning,
@@ -35,7 +35,7 @@ if TYPE_CHECKING:
 #
 
 
-def run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
+def run_host_op(state: "State", host: "Host", op_hash: str) -> bool:
     state.trigger_callbacks("operation_host_start", host, op_hash)
 
     if op_hash not in state.ops[host]:
@@ -59,7 +59,7 @@ def run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
             host.executing_op_hash = None
 
 
-def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
+def _run_host_op(state: "State", host: "Host", op_hash: str) -> bool:
     op_data = state.get_op_data_for_host(host, op_hash)
     global_arguments = op_data.global_arguments
 
@@ -104,6 +104,8 @@ def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
             if isinstance(command, FunctionCommand):
                 try:
                     status = command.execute(state, host, connector_arguments)
+                except NestedOperationError:
+                    host.log_styled("Error in nested operation", fg="red", log_func=logger.error)
                 except Exception as e:
                     # Custom functions could do anything, so expect anything!
                     logger.warning(traceback.format_exc())
@@ -194,11 +196,11 @@ def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
         host_results.ops += 1
         host_results.success_ops += 1
 
-        _status_log = "Success" if executed_commands > 0 else "No changes"
+        _status_text = "Success" if executed_commands > 0 else "No changes"
         if retry_attempt > 0:
-            _status_log = f"{_status_log} on retry {retry_attempt}"
+            _status_text = f"{_status_text} on retry {retry_attempt}"
 
-        _click_log_status = click.style(_status_log, "green")
+        _click_log_status = click.style(_status_text, "green" if executed_commands > 0 else "cyan")
         logger.info("{0}{1}".format(host.print_prefix, _click_log_status))
 
         state.trigger_callbacks("operation_host_success", host, op_hash, retry_attempt)
@@ -278,7 +280,7 @@ def _run_serial_ops(state: "State"):
     Run all ops for all servers, one server at a time.
     """
 
-    for host in list(state.inventory.iter_active_hosts()):
+    for host in list(state.inventory.get_active_hosts()):
         host_operations = product([host], state.get_op_order())
         with progress_spinner(host_operations) as progress:
             try:
@@ -296,7 +298,7 @@ def _run_no_wait_ops(state: "State"):
     Run all ops for all servers at once.
     """
 
-    hosts_operations = product(state.inventory.iter_active_hosts(), state.get_op_order())
+    hosts_operations = product(state.inventory.get_active_hosts(), state.get_op_order())
     with progress_spinner(hosts_operations) as progress:
         # Spawn greenlet for each host to run *all* ops
         if state.pool is None:
@@ -308,7 +310,7 @@ def _run_no_wait_ops(state: "State"):
                 host,
                 progress=progress,
             )
-            for host in state.inventory.iter_active_hosts()
+            for host in state.inventory.get_active_hosts()
         ]
         gevent.joinall(greenlets)
 
@@ -326,9 +328,9 @@ def _run_single_op(state: "State", op_hash: str):
     failed_hosts = set()
 
     if op_meta.global_arguments["_serial"]:
-        with progress_spinner(state.inventory.iter_active_hosts()) as progress:
+        with progress_spinner(state.inventory.get_active_hosts()) as progress:
             # For each host, run the op
-            for host in state.inventory.iter_active_hosts():
+            for host in state.inventory.get_active_hosts():
                 result = _run_host_op_with_context(state, host, op_hash)
                 progress(host)
 
@@ -337,12 +339,12 @@ def _run_single_op(state: "State", op_hash: str):
 
     else:
         # Start with the whole inventory in one batch
-        batches = [list(state.inventory.iter_active_hosts())]
+        batches = [list(state.inventory.get_active_hosts())]
 
         # If parallel set break up the inventory into a series of batches
         parallel = op_meta.global_arguments["_parallel"]
         if parallel:
-            hosts = list(state.inventory.iter_active_hosts())
+            hosts = list(state.inventory.get_active_hosts())
             batches = [hosts[i : i + parallel] for i in range(0, len(hosts), parallel)]
 
         for batch in batches:

pyinfra/api/util.py

@@ -10,7 +10,7 @@ from socket import error as socket_error, timeout as timeout_error
 from typing import IO, TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
 
 import click
-from jinja2 import Environment, FileSystemLoader, StrictUndefined
+from jinja2 import Environment, FileSystemLoader, StrictUndefined, Template
 from paramiko import SSHException
 from typeguard import TypeCheckError, check_type
 
@@ -26,7 +26,7 @@ if TYPE_CHECKING:
 BLOCKSIZE = 65536
 
 # Caches
-TEMPLATES: Dict[Any, Any] = {}
+TEMPLATES: Dict[str, Template] = {}
 FILE_SHAS: Dict[Any, Any] = {}
 
 PYINFRA_INSTALL_DIR = path.normpath(path.join(path.dirname(__file__), ".."))
@@ -139,7 +139,9 @@ def get_operation_order_from_stack(state: "State"):
     return line_numbers
 
 
-def get_template(filename_or_io: str | IO, jinja_env_kwargs: dict[str, Any] | None = None):
+def get_template(
+    filename_or_io: str | IO, jinja_env_kwargs: dict[str, Any] | None = None
+) -> Template:
     """
     Gets a jinja2 ``Template`` object for the input filename or string, with caching
     based on the filename of the template, or the SHA1 of the input string.
@@ -155,10 +157,11 @@ def get_template(filename_or_io: str | IO, jinja_env_kwargs: dict[str, Any] | No
     with file_data as file_io:
         template_string = file_io.read()
 
+    default_loader = FileSystemLoader(getcwd())
     template = Environment(
         undefined=StrictUndefined,
         keep_trailing_newline=True,
-        loader=FileSystemLoader(getcwd()),
+        loader=jinja_env_kwargs.pop("loader", default_loader),
         **jinja_env_kwargs,
     ).from_string(template_string)
 
@@ -219,7 +222,11 @@ def log_operation_start(
 
 
 def log_error_or_warning(
-    host: "Host", ignore_errors: bool, description: str = "", continue_on_error: bool = False
+    host: "Host",
+    ignore_errors: bool,
+    description: str = "",
+    continue_on_error: bool = False,
+    exception: Exception | None = None,
 ) -> None:
     log_func = logger.error
     log_color = "red"
@@ -234,6 +241,16 @@ def log_error_or_warning(
         if description:
             log_text = f"{log_text}: "
 
+    if exception:
+        exc = exception.__cause__ or exception
+        exc_text = "{0}: {1}".format(type(exc).__name__, exc)
+        log_func(
+            "{0}{1}".format(
+                host.print_prefix,
+                click.style(exc_text, log_color),
+            ),
+        )
+
     log_func(
         "{0}{1}{2}".format(
             host.print_prefix,

pyinfra/connectors/docker.py

@@ -26,10 +26,14 @@ if TYPE_CHECKING:
 
 class ConnectorData(TypedDict):
     docker_identifier: str
+    docker_platform: str
+    docker_architecture: str
 
 
 connector_data_meta: dict[str, DataMeta] = {
     "docker_identifier": DataMeta("ID of container or image to start from"),
+    "docker_platform": DataMeta("Platform to use for Docker image (e.g., linux/amd64)"),
+    "docker_architecture": DataMeta("Architecture to use for Docker image (e.g., amd64, arm64)"),
 }
 
 
@@ -108,9 +112,29 @@ class DockerConnector(BaseConnector):
         return container_id, True
 
     def _start_docker_image(self, image_name):
+        docker_cmd_parts = [
+            self.docker_cmd,
+            "run",
+            "-d",
+        ]
+
+        if self.data.get("docker_platform"):
+            docker_cmd_parts.extend(["--platform", self.data["docker_platform"]])
+        if self.data.get("docker_architecture"):
+            docker_cmd_parts.extend(["--arch", self.data["docker_architecture"]])
+
+        docker_cmd_parts.extend(
+            [
+                image_name,
+                "tail",
+                "-f",
+                "/dev/null",
+            ]
+        )
+
         try:
             return local.shell(
-                f"{self.docker_cmd} run -d {image_name} tail -f /dev/null",
+                " ".join(docker_cmd_parts),
                 splitlines=True,
             )[-1]  # last line is the container ID
         except PyinfraError as e:

pyinfra/connectors/ssh.py

@@ -9,6 +9,7 @@ from typing import IO, TYPE_CHECKING, Any, Iterable, Optional, Protocol, Tuple
 
 import click
 from paramiko import AuthenticationException, BadHostKeyException, SFTPClient, SSHException
+from paramiko.agent import Agent
 from typing_extensions import TypedDict, Unpack, override
 
 from pyinfra import logger
@@ -286,10 +287,64 @@ class SSHConnector(BaseConnector):
                 f"Host key for {e.hostname} does not match.",
             )
 
+        except SSHException as e:
+            if self._retry_paramiko_agent_keys(hostname, kwargs, e):
+                return
+            raise
+
     @override
     def disconnect(self) -> None:
         self.get_file_transfer_connection.cache.clear()
 
+    def _retry_paramiko_agent_keys(
+        self,
+        hostname: str,
+        kwargs: dict[str, Any],
+        error: SSHException,
+    ) -> bool:
+        # Workaround for Paramiko multi-key bug (paramiko/paramiko#1390).
+        if "no existing session" not in str(error).lower():
+            return False
+
+        if not kwargs.get("allow_agent"):
+            return False
+
+        try:
+            agent_keys = list(Agent().get_keys())
+        except Exception:
+            return False
+
+        if not agent_keys:
+            return False
+
+        # Skip the first agent key, since Paramiko already attempted it
+        attempt_keys = agent_keys[1:] if len(agent_keys) > 1 else agent_keys
+
+        for agent_key in attempt_keys:
+            if self.client is not None:
+                try:
+                    self.client.close()
+                except Exception:
+                    pass
+
+            self.client = SSHClient()
+
+            single_key_kwargs = dict(kwargs)
+            single_key_kwargs["allow_agent"] = False
+            single_key_kwargs["pkey"] = agent_key
+
+            try:
+                self.client.connect(hostname, **single_key_kwargs)
+                return True
+            except AuthenticationException:
+                continue
+            except SSHException as retry_error:
+                if "no existing session" in str(retry_error).lower():
+                    continue
+                raise retry_error
+
+        return False
+
     @override
     def run_shell_command(
         self,
@@ -342,8 +397,10 @@ class SSHConnector(BaseConnector):
                 get_pty=_get_pty,
             )
 
+            # Write any stdin and then close it
             if _stdin:
                 write_stdin(_stdin, stdin_buffer)
+            stdin_buffer.close()
 
             combined_output = read_output_buffers(
                 stdout_buffer,