pyinfra 2.9.2__py2.py3-none-any.whl → 3.0__py2.py3-none-any.whl

pyinfra/api/operations.py

@@ -1,26 +1,27 @@
+from __future__ import annotations
+
 import traceback
 from itertools import product
 from socket import error as socket_error, timeout as timeout_error
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional, cast
 
 import click
 import gevent
 from paramiko import SSHException
 
-import pyinfra
 from pyinfra import logger
-from pyinfra.context import ctx_host
+from pyinfra.connectors.util import CommandOutput, OutputLine
+from pyinfra.context import ctx_host, ctx_state
 from pyinfra.progress import progress_spinner
 
-from .arguments import get_executor_kwarg_keys
+from .arguments import CONNECTOR_ARGUMENT_KEYS, ConnectorArguments
 from .command import FunctionCommand, PyinfraCommand, StringCommand
-from .exceptions import NoMoreHostsError, PyinfraError
+from .exceptions import PyinfraError
 from .util import (
     format_exception,
     log_error_or_warning,
     log_host_command_error,
     log_operation_start,
-    memoize,
     print_host_combined_output,
 )
 
@@ -29,192 +30,153 @@ 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):
+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(host, op_hash)
-    global_kwargs = op_data["global_kwargs"]
-
     op_meta = state.get_op_meta(op_hash)
+    logger.debug("Starting operation %r on %s", op_meta.names, host)
 
-    ignore_errors = global_kwargs["ignore_errors"]
-    continue_on_error = global_kwargs["continue_on_error"]
-
-    logger.debug("Starting operation %r on %s", op_meta["names"], host)
-
-    executor_kwarg_keys = get_executor_kwarg_keys()
-    base_executor_kwargs = {
-        key: global_kwargs[key] for key in executor_kwarg_keys if key in global_kwargs
-    }
-
-    def _run_shell_command(command, executor_kwargs):
-        status = False
-        combined_output_lines = []
-
-        try:
-            status, combined_output_lines = command.execute(state, host, executor_kwargs)
-        except (timeout_error, socket_error, SSHException) as e:
-            log_host_command_error(
-                host,
-                e,
-                timeout=global_kwargs["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_kwargs[condition_name]
-        if not condition_value:
-            return True
-
-        show_pre_or_post_condition_warning(condition_name)
+    if host.executing_op_hash is None:
+        host.executing_op_hash = op_hash
+    else:
+        host.nested_executing_op_hash = op_hash
 
-        _shell_command_status, _ = _run_shell_command(
-            StringCommand(condition_value),
-            base_executor_kwargs,
-        )
+    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 _shell_command_status:
-            return True
 
-        _log_msg = f"{condition_name} failed: {condition_value}"
-        log_error_or_warning(host, ignore_errors, description=_log_msg)
+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 ignore_errors:
-            return True
+    ignore_errors = global_arguments["_ignore_errors"]
+    continue_on_error = global_arguments["_continue_on_error"]
+    timeout = global_arguments.get("_timeout", 0)
 
-        state.trigger_callbacks("operation_host_error", host, op_hash)
-        return False
+    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
+    )
 
-    if not run_condition("precondition"):
-        return False
-
-    state.ops_run.add(op_hash)
-
-    if host.executing_op_hash is None:
-        host.executing_op_hash = op_hash
-    else:
-        host.nested_executing_op_hash = op_hash
-
-    return_status = False
     did_error = False
     executed_commands = 0
-    all_combined_output_lines = []
+    commands = []
+    all_output_lines: list[OutputLine] = []
 
-    for i, command in enumerate(op_data["commands"]):
-        status = False
+    for command in op_data.command_generator():
+        commands.append(command)
 
-        executor_kwargs = base_executor_kwargs.copy()
-        executor_kwargs.update(command.executor_kwargs)
+        status = False
 
-        # Now we attempt to execute the command
-        #
+        connector_arguments = base_connector_arguments.copy()
+        connector_arguments.update(command.connector_arguments)
 
         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, executor_kwargs)
-            except NoMoreHostsError:
-                status = False
-            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)
+                status = command.execute(state, host, connector_arguments)
+            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, executor_kwargs)
-            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, executor_kwargs)
+                status = command.execute(state, host, connector_arguments)
             except (timeout_error, socket_error, SSHException, IOError) as e:
-                _timeout = global_kwargs["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
-        state.results[host]["commands"] += 1
 
-    # Commands didn't break, so count our successes & return True!
-    else:
-        if not run_condition("postcondition"):
-            return False
+    # Handle results
+    #
 
-        if not did_error:
-            return_status = True
+    op_success = return_status = not did_error
+    host_results = state.get_results_for_host(host)
 
-    if return_status is True:
-        state.results[host]["ops"] += 1
-        state.results[host]["success_ops"] += 1
+    if did_error is False:
+        host_results.ops += 1
+        host_results.success_ops += 1
 
-        _status_log = "Success" if len(op_data["commands"]) > 0 else "No changes"
+        _status_log = "Success" if executed_commands > 0 else "No changes"
         _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_kwargs["on_success"]:
-            global_kwargs["on_success"](state, host, op_hash)
-
         state.trigger_callbacks("operation_host_success", host, op_hash)
     else:
         if ignore_errors:
-            state.results[host]["ignored_error_ops"] += 1
+            host_results.ignored_error_ops += 1
         else:
-            state.results[host]["error_ops"] += 1
+            host_results.error_ops += 1
 
         if executed_commands:
-            state.results[host]["partial_ops"] += 1
+            host_results.partial_ops += 1
 
-        _command_description = f"executed {executed_commands}/{len(op_data['commands'])} commands"
+        _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_kwargs["on_error"]:
-            global_kwargs["on_error"](state, host, op_hash)
-
         # Ignored, op "completes" w/ ignored error
         if ignore_errors:
-            state.results[host]["ops"] += 1
+            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),
+    )
 
-    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):
@@ -242,14 +204,11 @@ def _run_host_ops(state: "State", host: "Host", progress=None):
         if result is False:
             raise PyinfraError(
                 "Error in operation {0} on {1}".format(
-                    ", ".join(op_meta["names"]),
+                    ", ".join(op_meta.names),
                     host,
                 ),
             )
 
-        if pyinfra.is_cli:
-            click.echo(err=True)
-
 
 def _run_serial_ops(state: "State"):
     """
@@ -303,7 +262,7 @@ def _run_single_op(state: "State", op_hash: str):
 
     failed_hosts = set()
 
-    if op_meta["serial"]:
+    if op_meta.global_arguments["_serial"]:
         with progress_spinner(state.inventory.iter_active_hosts()) as progress:
             # For each host, run the op
             for host in state.inventory.iter_active_hosts():
@@ -318,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["parallel"]:
-            parallel = op_meta["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:
@@ -347,9 +305,6 @@ def _run_single_op(state: "State", op_hash: str):
     # Now all the batches/hosts are complete, fail any failures
     state.fail_hosts(failed_hosts)
 
-    if pyinfra.is_cli:
-        click.echo(err=True)
-
     state.trigger_callbacks("operation_end", op_hash)
 
 
@@ -366,15 +321,14 @@ def run_ops(state: "State", serial: bool = False, no_wait: bool = False):
     # Flag state as deploy in process
     state.is_executing = True
 
-    # Run all ops, but server by server
-    if serial:
-        _run_serial_ops(state)
-
-    # Run all the ops on each server in parallel (not waiting at each operation)
-    elif no_wait:
-        _run_no_wait_ops(state)
-
-    # Default: run all ops in order, waiting at each for all servers to complete
-    else:
-        for op_hash in state.get_op_order():
-            _run_single_op(state, op_hash)
+    with ctx_state.use(state):
+        # Run all ops, but server by server
+        if serial:
+            _run_serial_ops(state)
+        # Run all the ops on each server in parallel (not waiting at each operation)
+        elif no_wait:
+            _run_no_wait_ops(state)
+        # Default: run all ops in order, waiting at each for all servers to complete
+        else:
+            for op_hash in state.get_op_order():
+                _run_single_op(state, op_hash)