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

pyinfra/api/operations.py

@@ -10,11 +10,11 @@ import gevent
 from paramiko import SSHException
 
 from pyinfra import logger
-from pyinfra.connectors.util import CommandOutput
+from pyinfra.connectors.util import CommandOutput, OutputLine
 from pyinfra.context import ctx_host, ctx_state
 from pyinfra.progress import progress_spinner
 
-from .arguments import ConnectorArguments, get_connector_argument_keys
+from .arguments import CONNECTOR_ARGUMENT_KEYS, ConnectorArguments
 from .command import FunctionCommand, PyinfraCommand, StringCommand
 from .exceptions import PyinfraError
 from .util import (
@@ -22,7 +22,6 @@ from .util import (
     log_error_or_warning,
     log_host_command_error,
     log_operation_start,
-    memoize,
     print_host_combined_output,
 )
 
@@ -31,92 +30,53 @@ if TYPE_CHECKING:
     from .state import State
 
 
-@memoize
-def show_pre_or_post_condition_warning(condition_name):
-    logger.warning("The `{0}` argument is in beta!".format(condition_name))
+# Run a single host operation
+#
 
 
-def run_host_op(state: "State", host: "Host", op_hash) -> Optional[bool]:
+def run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
     state.trigger_callbacks("operation_host_start", host, op_hash)
 
     if op_hash not in state.ops[host]:
         logger.info("{0}{1}".format(host.print_prefix, click.style("Skipped", "blue")))
         return True
 
-    op_data = state.get_op_data_for_host(host, op_hash)
-    global_arguments = op_data.global_arguments
-
     op_meta = state.get_op_meta(op_hash)
-
-    ignore_errors = global_arguments["_ignore_errors"]
-    continue_on_error = global_arguments["_continue_on_error"]
-
     logger.debug("Starting operation %r on %s", op_meta.names, host)
 
-    executor_kwarg_keys = get_connector_argument_keys()
-    base_connector_arguments: ConnectorArguments = (
-        cast(  # https://github.com/python/mypy/issues/10371
-            ConnectorArguments,
-            {key: global_arguments[key] for key in executor_kwarg_keys if key in global_arguments},
-        )
-    )
-
-    def _run_shell_command(command, connector_arguments):
-        status = False
-        combined_output_lines = CommandOutput([])
-
-        try:
-            status, combined_output_lines = command.execute(state, host, connector_arguments)
-        except (timeout_error, socket_error, SSHException) as e:
-            log_host_command_error(
-                host,
-                e,
-                timeout=global_arguments["_timeout"],
-            )
-
-        # If we failed and have no already printed the stderr, print it
-        if status is False and not state.print_output:
-            print_host_combined_output(host, combined_output_lines)
-
-        return status, combined_output_lines
-
-    def run_condition(condition_name: str) -> bool:
-        condition_value = global_arguments[condition_name]
-        if not condition_value:
-            return True
-
-        show_pre_or_post_condition_warning(condition_name)
-
-        _shell_command_status, _ = _run_shell_command(
-            StringCommand(condition_value),
-            base_connector_arguments,
-        )
-
-        if _shell_command_status:
-            return True
+    if host.executing_op_hash is None:
+        host.executing_op_hash = op_hash
+    else:
+        host.nested_executing_op_hash = op_hash
 
-        _log_msg = f"{condition_name} failed: {condition_value}"
-        log_error_or_warning(host, ignore_errors, description=_log_msg)
+    try:
+        return _run_host_op(state, host, op_hash)
+    finally:
+        if host.nested_executing_op_hash:
+            host.nested_executing_op_hash = None
+        else:
+            host.executing_op_hash = None
 
-        if ignore_errors:
-            return True
 
-        state.trigger_callbacks("operation_host_error", host, op_hash)
-        return False
+def _run_host_op(state: "State", host: "Host", op_hash: str) -> Optional[bool]:
+    op_data = state.get_op_data_for_host(host, op_hash)
+    global_arguments = op_data.global_arguments
 
-    if not run_condition("_precondition"):
-        return False
+    ignore_errors = global_arguments["_ignore_errors"]
+    continue_on_error = global_arguments["_continue_on_error"]
+    timeout = global_arguments.get("_timeout", 0)
 
-    if host.executing_op_hash is None:
-        host.executing_op_hash = op_hash
-    else:
-        host.nested_executing_op_hash = op_hash
+    executor_kwarg_keys = CONNECTOR_ARGUMENT_KEYS
+    # See: https://github.com/python/mypy/issues/10371
+    base_connector_arguments: ConnectorArguments = cast(
+        ConnectorArguments,
+        {key: global_arguments[key] for key in executor_kwarg_keys if key in global_arguments},  # type: ignore[literal-required] # noqa
+    )
 
-    return_status: Optional[bool] = False
     did_error = False
     executed_commands = 0
     commands = []
-    all_combined_output_lines = []
+    all_output_lines: list[OutputLine] = []
 
     for command in op_data.command_generator():
         commands.append(command)
@@ -126,54 +86,58 @@ def run_host_op(state: "State", host: "Host", op_hash) -> Optional[bool]:
         connector_arguments = base_connector_arguments.copy()
         connector_arguments.update(command.connector_arguments)
 
-        # Now we attempt to execute the command
-        #
-
         if not isinstance(command, PyinfraCommand):
             raise TypeError("{0} is an invalid pyinfra command!".format(command))
 
         if isinstance(command, FunctionCommand):
             try:
                 status = command.execute(state, host, connector_arguments)
-            except Exception as e:  # Custom functions could do anything, so expect anything!
-                _formatted_exc = format_exception(e)
-                _error_msg = "Unexpected error in Python callback: {0}".format(_formatted_exc)
-                _error_msg_styled = click.style(_error_msg, "red")
-                _error_log = "{0}{1}".format(host.print_prefix, _error_msg_styled)
+            except Exception as e:
+                # Custom functions could do anything, so expect anything!
                 logger.warning(traceback.format_exc())
-                logger.error(_error_log)
+                host.log_styled(
+                    f"Unexpected error in Python callback: {format_exception(e)}",
+                    fg="red",
+                    log_func=logger.warning,
+                )
 
         elif isinstance(command, StringCommand):
-            status, combined_output_lines = _run_shell_command(command, connector_arguments)
-            all_combined_output_lines.extend(combined_output_lines)
+            output_lines = CommandOutput([])
+            try:
+                status, output_lines = command.execute(
+                    state,
+                    host,
+                    connector_arguments,
+                )
+            except (timeout_error, socket_error, SSHException) as e:
+                log_host_command_error(host, e, timeout=timeout)
+            all_output_lines.extend(output_lines)
+            # If we failed and have not already printed the stderr, print it
+            if status is False and not state.print_output:
+                print_host_combined_output(host, output_lines)
 
         else:
             try:
                 status = command.execute(state, host, connector_arguments)
             except (timeout_error, socket_error, SSHException, IOError) as e:
-                _timeout = global_arguments["_timeout"]
-                log_host_command_error(host, e, timeout=_timeout)
+                log_host_command_error(host, e, timeout=timeout)
 
         # Break the loop to trigger a failure
         if status is False:
+            did_error = True
             if continue_on_error is True:
-                did_error = True
                 continue
             break
 
         executed_commands += 1
 
-    # Commands didn't break, so count our successes & return True!
-    else:
-        if not run_condition("_postcondition"):
-            return False
-
-        if not did_error:
-            return_status = True
+    # Handle results
+    #
 
+    op_success = return_status = not did_error
     host_results = state.get_results_for_host(host)
 
-    if return_status is True:
+    if did_error is False:
         host_results.ops += 1
         host_results.success_ops += 1
 
@@ -181,10 +145,6 @@ def run_host_op(state: "State", host: "Host", op_hash) -> Optional[bool]:
         _click_log_status = click.style(_status_log, "green")
         logger.info("{0}{1}".format(host.print_prefix, _click_log_status))
 
-        # Trigger any success handler
-        if global_arguments["_on_success"]:
-            global_arguments["_on_success"](state, host, op_hash)
-
         state.trigger_callbacks("operation_host_success", host, op_hash)
     else:
         if ignore_errors:
@@ -198,31 +158,25 @@ def run_host_op(state: "State", host: "Host", op_hash) -> Optional[bool]:
         _command_description = f"executed {executed_commands} commands"
         log_error_or_warning(host, ignore_errors, _command_description, continue_on_error)
 
-        # Always trigger any error handler
-        if global_arguments["_on_error"]:
-            global_arguments["_on_error"](state, host, op_hash)
-
         # Ignored, op "completes" w/ ignored error
         if ignore_errors:
             host_results.ops += 1
+            return_status = True
 
         # Unignored error -> False
         state.trigger_callbacks("operation_host_error", host, op_hash)
 
-        if ignore_errors:
-            return_status = True
+    op_data.operation_meta.set_complete(
+        op_success,
+        commands,
+        CommandOutput(all_output_lines),
+    )
 
-    if return_status is False or (return_status is True and executed_commands > 0):
-        op_data.operation_meta.set_result(return_status)
-    op_data.operation_meta.set_commands(commands)
-    op_data.operation_meta.set_combined_output_lines(all_combined_output_lines)
+    return return_status
 
-    if host.nested_executing_op_hash:
-        host.nested_executing_op_hash = None
-    else:
-        host.executing_op_hash = None
 
-    return return_status
+# Run all operations strategies
+#
 
 
 def _run_host_op_with_context(state: "State", host: "Host", op_hash: str):
@@ -323,10 +277,9 @@ def _run_single_op(state: "State", op_hash: str):
         batches = [list(state.inventory.iter_active_hosts())]
 
         # If parallel set break up the inventory into a series of batches
-        if op_meta.global_arguments["_parallel"]:
-            parallel = op_meta.global_arguments["_parallel"]
+        parallel = op_meta.global_arguments["_parallel"]
+        if parallel:
             hosts = list(state.inventory.iter_active_hosts())
-
             batches = [hosts[i : i + parallel] for i in range(0, len(hosts), parallel)]
 
         for batch in batches:

pyinfra/api/state.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-from contextlib import contextmanager
+from collections import defaultdict
 from dataclasses import dataclass
+from enum import IntEnum
 from graphlib import CycleError, TopologicalSorter
 from multiprocessing import cpu_count
-from typing import TYPE_CHECKING, Iterator, Optional
-from uuid import uuid4
+from typing import TYPE_CHECKING, Callable, Iterator, Optional
 
 from gevent.pool import Pool
 from paramiko import PKey
@@ -14,7 +14,6 @@ from pyinfra import logger
 
 from .config import Config
 from .exceptions import PyinfraError
-from .util import sha1_hash
 
 if TYPE_CHECKING:
     from pyinfra.api.arguments import AllArguments
@@ -83,22 +82,35 @@ class BaseStateCallback:
         pass
 
 
+class StateStage(IntEnum):
+    # Setup - collect inventory & data
+    Setup = 1
+    # Connect - connect to the inventory
+    Connect = 2
+    # Prepare - detect operation changes
+    Prepare = 3
+    # Execute - execute operations
+    Execute = 4
+    # Disconnect - disconnect from the inventory
+    Disconnect = 5
+
+
 class StateOperationMeta:
     names: set[str]
     args: list[str]
-    op_order: tuple[int]
+    op_order: tuple[int, ...]
     global_arguments: "AllArguments"
 
-    def __init__(self, op_order: tuple[int]):
+    def __init__(self, op_order: tuple[int, ...]):
         self.op_order = op_order
         self.names = set()
         self.args = []
-        self.global_arguments = {}
+        self.global_arguments = {}  # type: ignore
 
 
 @dataclass
 class StateOperationHostData:
-    command_generator: Iterator["PyinfraCommand"]
+    command_generator: Callable[[], Iterator["PyinfraCommand"]]
     global_arguments: "AllArguments"
     operation_meta: "OperationMeta"
     parent_op_hash: Optional[str] = None
@@ -138,6 +150,11 @@ class State:
     # Main gevent pool
     pool: "Pool"
 
+    # Current stage this state is in
+    current_stage: StateStage = StateStage.Setup
+    # Warning counters by stage
+    stage_warnings: dict[StateStage, int] = defaultdict(int)
+
     # Whether we are executing operations (ie hosts are all ready)
     is_executing: bool = False
 
@@ -157,6 +174,7 @@ class State:
     current_deploy_filename: Optional[str] = None
     current_exec_filename: Optional[str] = None
     current_op_file_number: int = 0
+    should_raise_failed_hosts: Optional[Callable[["State"], bool]] = None
 
     def __init__(
         self,
@@ -165,7 +183,8 @@ class State:
         check_for_changes: bool = True,
         **kwargs,
     ):
-        """Initializes the state, the main Pyinfra
+        """
+        Initializes the state, the main Pyinfra
 
         Args:
             inventory (Optional[Inventory], optional): The inventory. Defaults to None.
@@ -254,6 +273,17 @@ class State:
 
         self.initialised = True
 
+    def set_stage(self, stage: StateStage) -> None:
+        if stage < self.current_stage:
+            raise Exception("State stage cannot go backwards!")
+        self.current_stage = stage
+
+    def increment_warning_counter(self) -> None:
+        self.stage_warnings[self.current_stage] += 1
+
+    def get_warning_counter(self) -> int:
+        return self.stage_warnings[self.current_stage]
+
     def should_check_for_changes(self):
         return self.check_for_changes
 
@@ -269,16 +299,6 @@ class State:
             func = getattr(handler, method_name)
             func(self, *args, **kwargs)
 
-    @contextmanager
-    def preserve_loop_order(self, items):
-        logger.warning(
-            (
-                "Using `state.preserve_loop_order` is not longer required for operations to be "
-                "executed in correct loop order and can be safely removed."
-            ),
-        )
-        yield lambda: items
-
     def get_op_order(self):
         ts: TopologicalSorter = TopologicalSorter()
 
@@ -324,10 +344,19 @@ class State:
     def get_results_for_host(self, host: "Host") -> StateHostResults:
         return self.results[host]
 
-    def get_op_data_for_host(self, host: "Host", op_hash: str):
+    def get_op_data_for_host(
+        self,
+        host: "Host",
+        op_hash: str,
+    ) -> StateOperationHostData:
         return self.ops[host][op_hash]
 
-    def set_op_data_for_host(self, host: "Host", op_hash: str, op_data: StateOperationHostData):
+    def set_op_data_for_host(
+        self,
+        host: "Host",
+        op_hash: str,
+        op_data: StateOperationHostData,
+    ):
         self.ops[host][op_hash] = op_data
 
     def activate_host(self, host: "Host"):
@@ -375,6 +404,9 @@ class State:
             percent_failed = (1 - len(active_hosts) / activated_count) * 100
 
             if percent_failed > self.config.FAIL_PERCENT:
+                if self.should_raise_failed_hosts and self.should_raise_failed_hosts(self) is False:
+                    return
+
                 raise PyinfraError(
                     "Over {0}% of hosts failed ({1}%)".format(
                         self.config.FAIL_PERCENT,
@@ -392,16 +424,3 @@ class State:
         if not isinstance(limit_hosts, list):
             return True
         return host in limit_hosts
-
-    def get_temp_filename(self, hash_key: Optional[str] = None, hash_filename: bool = True):
-        """
-        Generate a temporary filename for this deploy.
-        """
-
-        if not hash_key:
-            hash_key = str(uuid4())
-
-        if hash_filename:
-            hash_key = sha1_hash(hash_key)
-
-        return "{0}/pyinfra-{1}".format(self.config.TEMP_DIR, hash_key)

pyinfra/api/util.py

@@ -1,13 +1,17 @@
+from __future__ import annotations
+
 from functools import wraps
 from hashlib import sha1
 from inspect import getframeinfo, stack
+from io import BytesIO, StringIO
 from os import getcwd, path, stat
 from socket import error as socket_error, timeout as timeout_error
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from typing import IO, TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
 
 import click
 from jinja2 import Environment, FileSystemLoader, StrictUndefined
 from paramiko import SSHException
+from typeguard import TypeCheckError, check_type
 
 import pyinfra
 from pyinfra import logger
@@ -24,7 +28,7 @@ BLOCKSIZE = 65536
 TEMPLATES: Dict[Any, Any] = {}
 FILE_SHAS: Dict[Any, Any] = {}
 
-PYINFRA_API_DIR = path.dirname(__file__)
+PYINFRA_INSTALL_DIR = path.normpath(path.join(path.dirname(__file__), ".."))
 
 
 def get_file_path(state: "State", filename: str):
@@ -105,16 +109,16 @@ def get_caller_frameinfo(frame_offset: int = 0):
 
 
 def get_operation_order_from_stack(state: "State"):
+
     stack_items = list(reversed(stack()))
 
+    i = 0
     # Find the *first* occurrence of our deploy file in the reversed stack
     if state.current_deploy_filename:
         for i, stack_item in enumerate(stack_items):
             frame = getframeinfo(stack_item[0])
             if frame.filename == state.current_deploy_filename:
                 break
-    else:
-        i = 0
 
     # Now generate a list of line numbers *following that file*
     line_numbers = []
@@ -125,7 +129,7 @@ def get_operation_order_from_stack(state: "State"):
     for stack_item in stack_items[i:]:
         frame = getframeinfo(stack_item[0])
 
-        if frame.filename.startswith(PYINFRA_API_DIR):
+        if frame.filename.startswith(PYINFRA_INSTALL_DIR):
             continue
 
         line_numbers.append(frame.lineno)
@@ -135,7 +139,7 @@ def get_operation_order_from_stack(state: "State"):
     return line_numbers
 
 
-def get_template(filename_or_io: str):
+def get_template(filename_or_io: str | IO):
     """
     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.
@@ -237,7 +241,7 @@ def log_error_or_warning(
     )
 
 
-def log_host_command_error(host: "Host", e: Exception, timeout: int = 0) -> None:
+def log_host_command_error(host: "Host", e: Exception, timeout: int | None = 0) -> None:
     if isinstance(e, timeout_error):
         logger.error(
             "{0}{1}".format(
@@ -297,19 +301,27 @@ def make_hash(obj):
             if isinstance(obj, int)
             # Constants - the values can change between hosts but we should still
             # group them under the same operation hash.
-            else "_PYINFRA_CONSTANT"
-            if obj in (True, False, None)
-            # Plain strings
-            else obj
-            if isinstance(obj, str)
-            # Objects with __name__s
-            else obj.__name__
-            if hasattr(obj, "__name__")
-            # Objects with names
-            else obj.name
-            if hasattr(obj, "name")
-            # Repr anything else
-            else repr(obj)
+            else (
+                "_PYINFRA_CONSTANT"
+                if obj in (True, False, None)
+                # Plain strings
+                else (
+                    obj
+                    if isinstance(obj, str)
+                    # Objects with __name__s
+                    else (
+                        obj.__name__
+                        if hasattr(obj, "__name__")
+                        # Objects with names
+                        else (
+                            obj.name
+                            if hasattr(obj, "name")
+                            # Repr anything else
+                            else repr(obj)
+                        )
+                    )
+                )
+            )
         )
 
     return sha1_hash(hash_string)
@@ -321,7 +333,11 @@ class get_file_io:
     will open and close filenames, and leave IO objects alone.
     """
 
-    close = False
+    filename_or_io: Union[str, IO[Any]]
+    mode: str
+
+    _close: bool = False
+    _file_io: IO[Any]
 
     def __init__(self, filename_or_io, mode="rb"):
         if not (
@@ -336,29 +352,32 @@ class get_file_io:
                 ),
             )
 
+        # Convert any StringIO/BytesIO to the other to match the desired mode
+        if isinstance(filename_or_io, StringIO) and mode == "rb":
+            filename_or_io.seek(0)
+            filename_or_io = BytesIO(filename_or_io.read().encode())
+        if isinstance(filename_or_io, BytesIO) and mode == "r":
+            filename_or_io.seek(0)
+            filename_or_io = StringIO(filename_or_io.read().decode())
+
         self.filename_or_io = filename_or_io
         self.mode = mode
 
     def __enter__(self):
-        # If we have a read attribute, just use the object as-is
-        if hasattr(self.filename_or_io, "read"):
-            file_io = self.filename_or_io
-
-        # Otherwise, assume a filename and open it up
-        else:
+        if isinstance(self.filename_or_io, str):
             file_io = open(self.filename_or_io, self.mode)
-
-            # Attach to self for closing on __exit__
-            self.file_io = file_io
-            self.close = True
+            self._file_io = file_io
+            self._close = True
+        else:
+            file_io = self.filename_or_io
 
         # Ensure we're at the start of the file
         file_io.seek(0)
         return file_io
 
     def __exit__(self, type, value, traceback):
-        if self.close:
-            self.file_io.close()
+        if self._close:
+            self._file_io.close()
 
     @property
     def cache_key(self):
@@ -405,3 +424,15 @@ def get_path_permissions_mode(pathname: str):
 
     mode_octal = oct(stat(pathname).st_mode)
     return mode_octal[-3:]
+
+
+def raise_if_bad_type(
+    value: Any,
+    type_: Type,
+    exception: type[Exception],
+    message_prefix: str,
+):
+    try:
+        check_type(value, type_)
+    except TypeCheckError as e:
+        raise exception(f"{message_prefix}: {e}")