pyinfra 3.0.dev0__py2.py3-none-any.whl → 3.0.1__py2.py3-none-any.whl

pyinfra/api/host.py

@@ -1,9 +1,22 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Callable, Generator, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generator,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+from uuid import uuid4
 
 import click
+from typing_extensions import Unpack
 
 from pyinfra import logger
 from pyinfra.connectors.base import BaseConnector
@@ -11,7 +24,8 @@ from pyinfra.connectors.util import CommandOutput, remove_any_sudo_askpass_file
 
 from .connectors import get_execution_connector
 from .exceptions import ConnectError
-from .facts import get_host_fact
+from .facts import FactBase, ShortFactBase, get_host_fact
+from .util import memoize, sha1_hash
 
 if TYPE_CHECKING:
     from pyinfra.api.arguments import AllArguments
@@ -19,7 +33,9 @@ if TYPE_CHECKING:
     from pyinfra.api.state import State
 
 
-def extract_callable_datas(datas: list[Union[Callable[..., Any], Any]]) -> Generator[Any, Any, Any]:
+def extract_callable_datas(
+    datas: list[Union[Callable[..., Any], Any]],
+) -> Generator[Any, Any, Any]:
     for data in datas:
         # Support for dynamic data, ie @deploy wrapped data defaults where
         # the data is stored on the state temporarily.
@@ -94,17 +110,22 @@ class Host:
     in_op: bool = False
     in_callback_op: bool = False
     current_op_hash: Optional[str] = None
-    current_op_global_arguments: Optional["AllArguments"]
+    current_op_global_arguments: Optional["AllArguments"] = None
 
-    # Current context inside a @deploy function (op gen stage)
+    # Current context inside a @deploy function which become part of the op data
     in_deploy: bool = False
     current_deploy_name: Optional[str] = None
     current_deploy_kwargs = None
-    current_deploy_data = None
+
+    # @deploy decorator data is a bit different - we need to handle the case
+    # where we're evaluating an operation at runtime (current_op_) but also
+    # when ordering operations (current_) outside of an operation context.
+    current_op_deploy_data: Optional[dict[str, Any]] = None
+    current_deploy_data: Optional[dict[str, Any]] = None
 
     # Current context during operation execution
-    executing_op_hash = None
-    nested_executing_op_hash = None
+    executing_op_hash: Optional[str] = None
+    nested_executing_op_hash: Optional[str] = None
 
     loop_position: list[int]
 
@@ -133,7 +154,6 @@ class Host:
         self.loop_position = []
 
         self.connector_data = {}
-        self.current_op_global_arguments = {}
 
         # Append only list of operation hashes as called on this host, used to
         # generate a DAG to create the final operation order.
@@ -195,11 +215,15 @@ class Host:
             self.print_prefix_padding,
         )
 
-    def get_deploy_data(self):
-        if self.current_deploy_data:
-            return self.current_deploy_data
+    def log(self, message, log_func=logger.info):
+        log_func(f"{self.print_prefix}{message}")
 
-        return {}
+    def log_styled(self, message, log_func=logger.info, **kwargs):
+        message_styled = click.style(message, **kwargs)
+        self.log(message_styled, log_func=log_func)
+
+    def get_deploy_data(self):
+        return self.current_op_deploy_data or self.current_deploy_data or {}
 
     def noop(self, description):
         """
@@ -209,8 +233,25 @@ class Host:
         handler = logger.info if self.state.print_noop_info else logger.debug
         handler("{0}noop: {1}".format(self.print_prefix, description))
 
+    def when(self, condition: Callable[[], bool]):
+        return self.deploy(
+            "",
+            cast("AllArguments", {"_if": [condition]}),
+            {},
+            in_deploy=False,
+        )
+
+    def arguments(self, **arguments: Unpack["AllArguments"]):
+        return self.deploy("", arguments, {}, in_deploy=False)
+
     @contextmanager
-    def deploy(self, name: str, kwargs, data, in_deploy: bool = True):
+    def deploy(
+        self,
+        name: str,
+        kwargs: Optional["AllArguments"],
+        data: Optional[dict],
+        in_deploy: bool = True,
+    ):
         """
         Wraps a group of operations as a deploy, this should not be used
         directly, instead use ``pyinfra.api.deploy.deploy``.
@@ -227,6 +268,13 @@ class Host:
         old_deploy_data = self.current_deploy_data
         self.in_deploy = in_deploy
 
+        # Combine any old _ifs with the new ones
+        if old_deploy_kwargs and kwargs:
+            old_ifs = old_deploy_kwargs["_if"]
+            new_ifs = kwargs["_if"]
+            if old_ifs and new_ifs:
+                kwargs["_if"] = old_ifs + new_ifs
+
         # Set the new values
         self.current_deploy_name = name
         self.current_deploy_kwargs = kwargs
@@ -253,9 +301,48 @@ class Host:
             old_deploy_data,
         )
 
+    @memoize
+    def _get_temp_directory(self):
+        temp_directory = self.state.config.TEMP_DIR
+
+        if temp_directory is None:
+            # Unfortunate, but very hard to avoid, circular dependency, this method is memoized so
+            # performance isn't a concern.
+            from pyinfra.facts.server import TmpDir
+
+            temp_directory = self.get_fact(TmpDir)
+
+        if not temp_directory:
+            temp_directory = self.state.config.DEFAULT_TEMP_DIR
+
+        return temp_directory
+
+    def get_temp_filename(self, hash_key: Optional[str] = None, hash_filename: bool = True):
+        """
+        Generate a temporary filename for this deploy.
+        """
+
+        temp_directory = self._get_temp_directory()
+
+        if not hash_key:
+            hash_key = str(uuid4())
+
+        if hash_filename:
+            hash_key = sha1_hash(hash_key)
+
+        return "{0}/pyinfra-{1}".format(temp_directory, hash_key)
+
     # Host facts
     #
 
+    T = TypeVar("T")
+
+    @overload
+    def get_fact(self, name_or_cls: Type[FactBase[T]], *args, **kwargs) -> T: ...
+
+    @overload
+    def get_fact(self, name_or_cls: Type[ShortFactBase[T]], *args, **kwargs) -> T: ...
+
     def get_fact(self, name_or_cls, *args, **kwargs):
         """
         Get a fact for this host, reading from the cache if present.
@@ -323,7 +410,7 @@ class Host:
 
         self.state.trigger_callbacks("host_disconnect", self)
 
-    def run_shell_command(self, *args, **kwargs) -> Tuple[bool, CommandOutput]:
+    def run_shell_command(self, *args, **kwargs) -> tuple[bool, CommandOutput]:
         """
         Low level method to execute a shell command on the host via it's configured connector.
         """

pyinfra/api/inventory.py

@@ -37,6 +37,10 @@ class Inventory:
 
     state: "State"
 
+    @staticmethod
+    def empty():
+        return Inventory(([], {}))
+
     def __init__(self, names_data, override_data=None, **groups):
         # Setup basics
         self.groups = defaultdict(list)  # lists of Host objects

pyinfra/api/operation.py

@@ -8,15 +8,19 @@ to the deploy state. This is then run later by pyinfra's ``__main__`` or the
 from __future__ import annotations
 
 from functools import wraps
+from inspect import signature
 from io import StringIO
 from types import FunctionType
-from typing import Any, Iterator, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Generator, Iterator, Optional, cast
+
+from typing_extensions import ParamSpec
 
 import pyinfra
 from pyinfra import context, logger
 from pyinfra.context import ctx_host, ctx_state
 
-from .arguments import AllArguments, get_execution_kwarg_keys, pop_global_arguments
+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 .host import Host
@@ -32,53 +36,115 @@ from .util import (
 
 op_meta_default = object()
 
+if TYPE_CHECKING:
+    from pyinfra.connectors.util import CommandOutput
+
 
 class OperationMeta:
-    combined_output_lines = None
-    commands: Optional[list[Any]] = None
-    changed: bool = False
-    success: Optional[bool] = None
+    _hash: str
+
+    _combined_output: Optional["CommandOutput"] = None
+    _commands: Optional[list[Any]] = None
+    _maybe_is_change: Optional[bool] = None
+    _success: Optional[bool] = None
 
-    def __init__(self, hash=None, is_change=False):
-        self.hash = hash
-        self.changed = is_change
+    def __init__(self, hash, is_change: Optional[bool]):
+        self._hash = hash
+        self._maybe_is_change = is_change
 
     def __repr__(self) -> str:
         """
         Return Operation object as a string.
         """
 
-        return f"OperationMeta(changed={self.changed}, hash={self.hash})"
+        if self._commands is not None:
+            return (
+                "OperationMeta(executed=True, "
+                f"success={self.did_succeed}, hash={self._hash}, commands={len(self._commands)})"
+            )
+        return (
+            "OperationMeta(executed=False, "
+            f"maybeChange={self._maybe_is_change}, hash={self._hash})"
+        )
+
+    # Completion & status checks
+    def set_complete(
+        self,
+        success: bool,
+        commands: list[Any],
+        combined_output: "CommandOutput",
+    ) -> None:
+        if self.is_complete():
+            raise RuntimeError("Cannot complete an already complete operation")
+        self._success = success
+        self._commands = commands
+        self._combined_output = combined_output
+
+    def is_complete(self) -> bool:
+        return self._success is not None
+
+    def _raise_if_not_complete(self) -> None:
+        if not self.is_complete():
+            raise RuntimeError("Cannot evaluate operation result before execution")
 
-    def set_combined_output_lines(self, combined_output_lines):
-        self.combined_output_lines = combined_output_lines
+    @property
+    def will_change(self) -> bool:
+        if self._maybe_is_change is not None:
+            return self._maybe_is_change
+
+        op_data = context.state.get_op_data_for_host(context.host, self._hash)
+        cmd_gen = op_data.command_generator
+        for _ in cmd_gen():
+            self._maybe_is_change = True
+            return True
+        self._maybe_is_change = False
+        return False
+
+    def _did_change(self) -> bool:
+        return bool(self._success and len(self._commands or []) > 0)
 
-    def set_commands(self, commands) -> None:
-        self.commands = commands
+    @property
+    def did_change(self):
+        return context.host.when(self._did_change)
 
-    def set_result(self, success: bool) -> None:
-        self.success = success
+    @property
+    def did_not_change(self):
+        return context.host.when(lambda: not self._did_change())
 
-    def _get_lines(self, types=("stdout", "stderr")):
-        if self.combined_output_lines is None:
-            raise AttributeError("Output is not available until operations have been executed")
+    def did_succeed(self, _raise_if_not_complete=True) -> bool:
+        if _raise_if_not_complete:
+            self._raise_if_not_complete()
+        return self._success is True
 
-        return [line for type_, line in self.combined_output_lines if type_ in types]
+    def did_error(self) -> bool:
+        self._raise_if_not_complete()
+        return self._success is False
 
+    # TODO: deprecated, remove in v4
     @property
-    def stdout_lines(self):
-        return self._get_lines(types=("stdout",))
+    def changed(self) -> bool:
+        if self.is_complete():
+            return self._did_change()
+        return self.will_change
 
     @property
-    def stderr_lines(self):
-        return self._get_lines(types=("stderr",))
+    def stdout_lines(self) -> list[str]:
+        self._raise_if_not_complete()
+        assert self._combined_output is not None
+        return self._combined_output.stdout_lines
 
     @property
-    def stdout(self):
+    def stderr_lines(self) -> list[str]:
+        self._raise_if_not_complete()
+        assert self._combined_output is not None
+        return self._combined_output.stderr_lines
+
+    @property
+    def stdout(self) -> str:
         return "\n".join(self.stdout_lines)
 
     @property
-    def stderr(self):
+    def stderr(self) -> str:
         return "\n".join(self.stderr_lines)
 
 
@@ -113,57 +179,67 @@ def add_op(state: State, op_func, *args, **kwargs):
     return results
 
 
+P = ParamSpec("P")
+
+
 def operation(
-    pipeline_facts=None,
     is_idempotent: bool = True,
-    idempotent_notice=None,
-):
+    idempotent_notice: Optional[str] = None,
+    is_deprecated: bool = False,
+    deprecated_for: Optional[str] = None,
+    _set_in_op: bool = True,
+) -> Callable[[Callable[P, Generator]], PyinfraOperation[P]]:
     """
     Decorator that takes a simple module function and turn it into the internal
     operation representation that consists of a list of commands + options
     (sudo, (sudo|su)_user, env).
     """
 
-    def decorator(f):
-        f.pipeline_facts = pipeline_facts
-        f.is_idempotent = is_idempotent
-        f.idempotent_notice = idempotent_notice
-        return _wrap_operation(f)
+    def decorator(f: Callable[P, Generator]) -> PyinfraOperation[P]:
+        f.is_idempotent = is_idempotent  # type: ignore[attr-defined]
+        f.idempotent_notice = idempotent_notice  # type: ignore[attr-defined]
+        f.is_deprecated = is_deprecated  # type: ignore[attr-defined]
+        f.deprecated_for = deprecated_for  # type: ignore[attr-defined]
+        return _wrap_operation(f, _set_in_op=_set_in_op)
 
     return decorator
 
 
-def _wrap_operation(func):
+def _wrap_operation(func: Callable[P, Generator], _set_in_op: bool = True) -> PyinfraOperation[P]:
     @wraps(func)
-    def decorated_func(*args, **kwargs):
+    def decorated_func(*args: P.args, **kwargs: P.kwargs) -> OperationMeta:
         state = context.state
         host = context.host
 
+        if host.in_op:
+            raise Exception(
+                "Operation called within another operation, this is not allowed! Use the `_inner` "
+                + "function to call the underlying operation."
+            )
+
+        if func.is_deprecated:  # type: ignore[attr-defined]
+            if func.deprecated_for:  # type: ignore[attr-defined]
+                logger.warning(
+                    f"The {get_operation_name_from_func(func)} operation is "
+                    + f"deprecated, please use: {func.deprecated_for}",  # type: ignore[attr-defined] # noqa
+                )
+            else:
+                logger.warning(f"The {get_operation_name_from_func(func)} operation is deprecated")
+
         # Configure operation
         #
         # Get the meta kwargs (globals that apply to all hosts)
         global_arguments, global_argument_keys = pop_global_arguments(kwargs)
 
-        # If this op is being called inside another, just return here
-        # (any unwanted/op-related kwargs removed above).
-        if host.in_op and not host.in_callback_op:
-            if global_argument_keys:
-                _error_msg = "Nested operation called with global arguments: {0} ({1})".format(
-                    global_argument_keys,
-                    get_call_location(),
-                )
-                raise PyinfraError(_error_msg)
-            return func(*args, **kwargs) or []
-
-        names, add_args = _generate_operation_name(func, host, kwargs, global_arguments)
-        op_order, op_hash = _solve_operation_consistency(names, state, host)
+        names, add_args = generate_operation_name(func, host, kwargs, global_arguments)
+        op_order, op_hash = solve_operation_consistency(names, state, host)
 
         # Ensure shared (between servers) operation meta, mutates state
-        op_meta = _ensure_shared_op_meta(state, op_hash, op_order, global_arguments, names)
+        op_meta = ensure_shared_op_meta(state, op_hash, op_order, global_arguments, names)
 
         # Attach normal args, if we're auto-naming this operation
         if add_args:
-            op_meta = _attach_args(op_meta, args, kwargs)
+            op_meta = attach_args(op_meta, args, kwargs)
 
         # Check if we're actually running the operation on this host
         # Run once and we've already added meta for this op? Stop here.
@@ -175,20 +251,30 @@ def _wrap_operation(func):
                     break
 
             if has_run:
-                return OperationMeta(op_hash)
+                return OperationMeta(op_hash, is_change=False)
+
+        # Grab a reference to any *current* deploy data as this may change when
+        # we later evaluate the operation at runtime.This means we put back the
+        # expected deploy data.
+        current_deploy_data = host.current_deploy_data
 
         # "Run" operation - here we make a generator that will yield out actual commands to execute
         # and, if we're diff-ing, we then iterate the generator now to determine if any changes
         # *would* be made based on the *current* remote state.
 
         def command_generator() -> Iterator[PyinfraCommand]:
-            host.in_op = True
-            # MY EYES, this is evil
-            host.in_callback_op = (
-                func.__name__ == "call" and func.__module__ == "pyinfra.operations.python"
-            )
+            # Check global _if argument function and do nothing if returns False
+            if state.is_executing:
+                _ifs = global_arguments.get("_if")
+                if isinstance(_ifs, list) and not all(_if() for _if in _ifs):
+                    return
+                elif callable(_ifs) and not _ifs():
+                    return
+
+            host.in_op = _set_in_op
             host.current_op_hash = op_hash
             host.current_op_global_arguments = global_arguments
+            host.current_op_deploy_data = current_deploy_data
 
             try:
                 for command in func(*args, **kwargs):
@@ -199,35 +285,54 @@ def _wrap_operation(func):
                 host.in_op = False
                 host.current_op_hash = None
                 host.current_op_global_arguments = None
+                host.current_op_deploy_data = None
 
-        op_is_change = False
+        op_is_change = None
         if state.should_check_for_changes():
-            for command in command_generator():
+            op_is_change = False
+            for _ in command_generator():
                 op_is_change = True
                 break
+        else:
+            # If not calling the op function to check for change we still want to ensure the args
+            # are valid, so use Signature.bind to trigger any TypeError.
+            signature(func).bind(*args, **kwargs)
 
         # Add host-specific operation data to state, this mutates state
-        operation_meta = _add_host_op_to_state(
-            state,
-            host,
-            op_hash,
-            op_is_change,
-            command_generator,
-            global_arguments,
-        )
+        host_meta = state.get_meta_for_host(host)
+        host_meta.ops += 1
+        if op_is_change:
+            host_meta.ops_change += 1
+        else:
+            host_meta.ops_no_change += 1
+
+        operation_meta = OperationMeta(op_hash, op_is_change)
+
+        # Add the server-relevant commands
+        op_data = StateOperationHostData(command_generator, global_arguments, operation_meta)
+        state.set_op_data_for_host(host, op_hash, op_data)
 
         # If we're already in the execution phase, execute this operation immediately
         if state.is_executing:
-            _execute_immediately(state, host, op_hash)
+            execute_immediately(state, host, op_hash)
 
         # Return result meta for use in deploy scripts
         return operation_meta
 
-    decorated_func._pyinfra_op = func  # type: ignore
-    return decorated_func
+    decorated_func._inner = func  # type: ignore[attr-defined]
+    return cast(PyinfraOperation[P], decorated_func)
 
 
-def _generate_operation_name(func, host, kwargs, global_arguments):
+def get_operation_name_from_func(func):
+    if func.__module__:
+        module_bits = func.__module__.split(".")
+        module_name = module_bits[-1]
+        return "{0}.{1}".format(module_name, func.__name__)
+    else:
+        return func.__name__
+
+
+def generate_operation_name(func, host, kwargs, global_arguments):
     # Generate an operation name if needed (Module/Operation format)
     name = global_arguments.get("name")
     add_args = False
@@ -235,14 +340,7 @@ def _generate_operation_name(func, host, kwargs, global_arguments):
         names = {name}
     else:
         add_args = True
-
-        if func.__module__:
-            module_bits = func.__module__.split(".")
-            module_name = module_bits[-1]
-            name = "{0}/{1}".format(module_name.title(), func.__name__.title())
-        else:
-            name = func.__name__
-
+        name = get_operation_name_from_func(func)
         names = {name}
 
     if host.current_deploy_name:
@@ -251,7 +349,7 @@ def _generate_operation_name(func, host, kwargs, global_arguments):
     return names, add_args
 
 
-def _solve_operation_consistency(names, state, host):
+def solve_operation_consistency(names, state, host):
     # Operation order is used to tie-break available nodes in the operation DAG, in CLI mode
     # we use stack call order so this matches as defined by the user deploy code.
     if pyinfra.is_cli:
@@ -286,16 +384,16 @@ def _solve_operation_consistency(names, state, host):
 
 
 # NOTE: this function mutates state.op_meta for this hash
-def _ensure_shared_op_meta(
+def ensure_shared_op_meta(
     state: State,
     op_hash: str,
-    op_order: Tuple[int],
+    op_order: tuple[int, ...],
     global_arguments: AllArguments,
-    names: Set[str],
+    names: set[str],
 ):
     op_meta = state.op_meta.setdefault(op_hash, StateOperationMeta(op_order))
 
-    for key in get_execution_kwarg_keys():
+    for key in EXECUTION_KWARG_KEYS:
         global_value = global_arguments.pop(key)  # type: ignore[misc]
         op_meta_value = op_meta.global_arguments.get(key, op_meta_default)
 
@@ -310,12 +408,7 @@ def _ensure_shared_op_meta(
     return op_meta
 
 
-def _execute_immediately(state, host, op_hash):
-    logger.warning(
-        f"Note: nested operations are currently in beta ({get_call_location()})\n"
-        "    More information: "
-        "https://docs.pyinfra.com/en/2.x/using-operations.html#nested-operations",
-    )
+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
@@ -331,7 +424,7 @@ def _get_arg_value(arg):
     return arg
 
 
-def _attach_args(op_meta, args, kwargs):
+def attach_args(op_meta, args, kwargs):
     for arg in args:
         if arg not in op_meta.args:
             op_meta.args.append(str(_get_arg_value(arg)))
@@ -343,30 +436,3 @@ def _attach_args(op_meta, args, kwargs):
             op_meta.args.append(arg)
 
     return op_meta
-
-
-# NOTE: this function mutates state.meta for this host
-def _add_host_op_to_state(
-    state: State,
-    host: Host,
-    op_hash: str,
-    is_change: bool,
-    command_generator,
-    global_arguments,
-) -> OperationMeta:
-    host_meta = state.get_meta_for_host(host)
-
-    host_meta.ops += 1
-
-    if is_change:
-        host_meta.ops_change += 1
-    else:
-        host_meta.ops_no_change += 1
-
-    operation_meta = OperationMeta(op_hash, is_change)
-
-    # Add the server-relevant commands
-    op_data = StateOperationHostData(command_generator, global_arguments, operation_meta)
-    state.set_op_data_for_host(host, op_hash, op_data)
-
-    return operation_meta