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

pyinfra/api/state.py

@@ -1,8 +1,11 @@
-from contextlib import contextmanager
+from __future__ import annotations
+
+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, Any, Dict, List, Optional, Set
-from uuid import uuid4
+from typing import TYPE_CHECKING, Callable, Iterator, Optional
 
 from gevent.pool import Pool
 from paramiko import PKey
@@ -10,12 +13,14 @@ from paramiko import PKey
 from pyinfra import logger
 
 from .config import Config
-from .exceptions import NoMoreHostsError, PyinfraError
-from .util import sha1_hash
+from .exceptions import PyinfraError
 
 if TYPE_CHECKING:
+    from pyinfra.api.arguments import AllArguments
+    from pyinfra.api.command import PyinfraCommand
     from pyinfra.api.host import Host
     from pyinfra.api.inventory import Inventory
+    from pyinfra.api.operation import OperationMeta
 
 
 # Work out the max parallel we can achieve with the open files limit of the user/process,
@@ -77,6 +82,58 @@ 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, ...]
+    global_arguments: "AllArguments"
+
+    def __init__(self, op_order: tuple[int, ...]):
+        self.op_order = op_order
+        self.names = set()
+        self.args = []
+        self.global_arguments = {}  # type: ignore
+
+
+@dataclass
+class StateOperationHostData:
+    command_generator: Callable[[], Iterator["PyinfraCommand"]]
+    global_arguments: "AllArguments"
+    operation_meta: "OperationMeta"
+    parent_op_hash: Optional[str] = None
+
+
+class StateHostMeta:
+    ops = 0
+    ops_change = 0
+    ops_no_change = 0
+    op_hashes: set[str]
+
+    def __init__(self):
+        self.op_hashes = set()
+
+
+class StateHostResults:
+    ops = 0
+    success_ops = 0
+    error_ops = 0
+    ignored_error_ops = 0
+    partial_ops = 0
+
+
 class State:
     """
     Manages state for a pyinfra deploy.
@@ -93,9 +150,18 @@ 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
 
+    # Whether we should check for operation changes as part of the operation ordering phase, this
+    # allows us to guesstimate which ops will result in changes on which hosts.
+    check_for_changes: bool = True
+
     print_noop_info: bool = False  # print "[host] noop: reason for noop"
     print_fact_info: bool = False  # print "loaded fact X"
     print_input: bool = False
@@ -108,20 +174,33 @@ 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, inventory: Optional["Inventory"] = None, config: Optional["Config"] = None, **kwargs
+        self,
+        inventory: Optional["Inventory"] = None,
+        config: Optional["Config"] = None,
+        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.
             config (Optional[Config], optional): The config object. Defaults to None.
         """
+        self.check_for_changes = check_for_changes
+
         if inventory:
             self.init(inventory, config, **kwargs)
 
-    def init(self, inventory: "Inventory", config: Optional["Config"], initial_limit=None):
+    def init(
+        self,
+        inventory: "Inventory",
+        config: Optional["Config"],
+        initial_limit=None,
+    ):
         # Config validation
         #
 
@@ -150,79 +229,63 @@ class State:
         # Actually initialise the state object
         #
 
-        self.callback_handlers: List[BaseStateCallback] = []
+        self.callback_handlers: list[BaseStateCallback] = []
 
         # Setup greenlet pools
         self.pool = Pool(config.PARALLEL)
         self.fact_pool = Pool(config.PARALLEL)
 
         # Private keys
-        self.private_keys: Dict[str, PKey] = {}
+        self.private_keys: dict[str, PKey] = {}
 
         # Assign inventory/config
         self.inventory = inventory
         self.config = config
 
         # Hosts we've activated at any time
-        self.activated_hosts: Set["Host"] = set()
+        self.activated_hosts: set["Host"] = set()
         # Active hosts that *haven't* failed yet
-        self.active_hosts: Set["Host"] = set()
+        self.active_hosts: set["Host"] = set()
         # Hosts that have failed
-        self.failed_hosts: Set["Host"] = set()
+        self.failed_hosts: set["Host"] = set()
 
         # Limit hosts changes dynamically to limit operations to a subset of hosts
-        self.limit_hosts: List["Host"] = initial_limit
+        self.limit_hosts: list["Host"] = initial_limit
 
         # Op basics
-        self.op_meta: Dict[str, dict] = {}  # maps operation hash -> names/etc
-        self.ops_run: Set[str] = set()  # list of ops which have been started/run
+        self.op_meta: dict[str, StateOperationMeta] = {}  # maps operation hash -> names/etc
 
         # Op dict for each host
-        self.ops: Dict["Host", dict] = {host: {} for host in inventory}
-
-        # Facts dict for each host
-        self.facts: Dict["Host", Any] = {host: {} for host in inventory}
+        self.ops: dict["Host", dict[str, StateOperationHostData]] = {host: {} for host in inventory}
 
         # Meta dict for each host
-        self.meta = {
-            host: {
-                "ops": 0,  # one function call in a deploy file
-                "ops_change": 0,
-                "ops_no_change": 0,
-                "commands": 0,  # actual # of commands to run
-                "op_hashes": set(),
-            }
-            for host in inventory
-        }
+        self.meta: dict["Host", StateHostMeta] = {host: StateHostMeta() for host in inventory}
 
         # Results dict for each host
-        self.results = {
-            host: {
-                "ops": 0,  # success_ops + failed ops w/ignore_errors
-                "success_ops": 0,
-                "error_ops": 0,
-                "ignored_error_ops": 0,
-                "partial_ops": 0,  # operations that had an error, but did do something
-                "commands": 0,
-            }
-            for host in inventory
+        self.results: dict["Host", StateHostResults] = {
+            host: StateHostResults() for host in inventory
         }
 
         # Assign state back references to inventory & config
         inventory.state = config.state = self
         for host in inventory:
-            host.state = self
+            host.init(self)
 
         self.initialised = True
 
-    def to_dict(self):
-        return {
-            "op_order": self.get_op_order(),
-            "ops": self.ops,
-            "facts": self.facts,
-            "meta": self.meta,
-            "results": self.results,
-        }
+    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
 
     def add_callback_handler(self, handler):
         if not isinstance(handler, BaseStateCallback):
@@ -236,18 +299,8 @@ 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()
+        ts: TopologicalSorter = TopologicalSorter()
 
         for host in self.inventory:
             for i, op_hash in enumerate(host.op_hash_order):
@@ -275,20 +328,35 @@ class State:
             # dependency order we order them by line numbers.
             node_group = sorted(
                 ts.get_ready(),
-                key=lambda op_hash: self.op_meta[op_hash]["op_order"],
+                key=lambda op_hash: self.op_meta[op_hash].op_order,
             )
             ts.done(*node_group)
             final_op_order.extend(node_group)
 
         return final_op_order
 
-    def get_op_meta(self, op_hash: str):
+    def get_op_meta(self, op_hash: str) -> StateOperationMeta:
         return self.op_meta[op_hash]
 
-    def get_op_data(self, host: "Host", op_hash: str):
+    def get_meta_for_host(self, host: "Host") -> StateHostMeta:
+        return self.meta[host]
+
+    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,
+    ) -> StateOperationHostData:
         return self.ops[host][op_hash]
 
-    def set_op_data(self, host: "Host", op_hash: str, op_data):
+    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"):
@@ -330,13 +398,16 @@ class State:
 
         # No hosts left!
         if not active_hosts:
-            raise NoMoreHostsError("No hosts remaining!")
+            raise PyinfraError("No hosts remaining!")
 
         if self.config.FAIL_PERCENT is not None:
             percent_failed = (1 - len(active_hosts) / activated_count) * 100
 
             if percent_failed > self.config.FAIL_PERCENT:
-                raise NoMoreHostsError(
+                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,
                         int(round(percent_failed)),
@@ -353,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,21 +1,25 @@
+from __future__ import annotations
+
 from functools import wraps
 from hashlib import sha1
-from inspect import getframeinfo, getfullargspec, stack
+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, Tuple
+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
 
 if TYPE_CHECKING:
     from pyinfra.api.host import Host
-    from pyinfra.api.state import State
+    from pyinfra.api.state import State, StateOperationMeta
+    from pyinfra.connectors.util import CommandOutput
 
 # 64kb chunks
 BLOCKSIZE = 65536
@@ -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):
@@ -40,28 +44,6 @@ def get_file_path(state: "State", filename: str):
     return path.join(relative_to, filename)
 
 
-def get_args_kwargs_spec(func: Callable[..., Any]) -> Tuple[List, Dict]:
-    args: List[Any] = []
-    kwargs: Dict[Any, Any] = {}
-
-    argspec = getfullargspec(func)
-    if not argspec.args:
-        return args, kwargs
-
-    if argspec.defaults:
-        kwargs = dict(
-            zip(
-                argspec.args[-len(argspec.defaults) :],
-                argspec.defaults,
-            ),
-        )
-        args = argspec.args[: -len(argspec.defaults)]
-    else:
-        args = argspec.args
-
-    return args, kwargs
-
-
 def get_kwargs_str(kwargs: Dict[Any, Any]):
     if not kwargs:
         return ""
@@ -85,14 +67,14 @@ def memoize(func: Callable[..., Any]):
     @wraps(func)
     def wrapper(*args, **kwargs):
         key = "{0}{1}".format(args, kwargs)
-        if key in wrapper.cache:
-            return wrapper.cache[key]
+        if key in wrapper.cache:  # type: ignore[attr-defined]
+            return wrapper.cache[key]  # type: ignore[attr-defined]
 
         value = func(*args, **kwargs)
-        wrapper.cache[key] = value
+        wrapper.cache[key] = value  # type: ignore[attr-defined]
         return value
 
-    wrapper.cache = {}  # type: ignore
+    wrapper.cache = {}  # type: ignore[attr-defined]
     return wrapper
 
 
@@ -127,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 = []
@@ -147,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)
@@ -157,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.
@@ -184,7 +166,7 @@ def get_template(filename_or_io: str):
     return template
 
 
-def sha1_hash(string: str):
+def sha1_hash(string: str) -> str:
     """
     Return the SHA1 of the input string.
     """
@@ -194,38 +176,30 @@ def sha1_hash(string: str):
     return hasher.hexdigest()
 
 
-def format_exception(e):
-    return "{0}{1}".format(e.__class__.__name__, e.args)
+def format_exception(e: Exception) -> str:
+    return f"{e.__class__.__name__}{e.args}"
 
 
-def print_host_combined_output(host: "Host", combined_output_lines):
-    for type_, line in combined_output_lines:
-        if type_ == "stderr":
-            logger.error(
-                "{0}{1}".format(
-                    host.print_prefix,
-                    click.style(line, "red"),
-                ),
-            )
+def print_host_combined_output(host: "Host", output: "CommandOutput") -> None:
+    for line in output:
+        if line.buffer_name == "stderr":
+            logger.error(f"{host.print_prefix}{click.style(line.line, 'red')}")
         else:
-            logger.error(
-                "{0}{1}".format(
-                    host.print_prefix,
-                    line,
-                ),
-            )
+            logger.error(f"{host.print_prefix}{line.line}")
 
 
-def log_operation_start(op_meta: Dict, op_types: Optional[List] = None, prefix: str = "--> "):
+def log_operation_start(
+    op_meta: "StateOperationMeta", op_types: Optional[List] = None, prefix: str = "--> "
+) -> None:
     op_types = op_types or []
-    if op_meta["serial"]:
+    if op_meta.global_arguments["_serial"]:
         op_types.append("serial")
-    if op_meta["run_once"]:
+    if op_meta.global_arguments["_run_once"]:
         op_types.append("run once")
 
     args = ""
-    if op_meta["args"]:
-        args = "({0})".format(", ".join(str(arg) for arg in op_meta["args"]))
+    if op_meta.args:
+        args = "({0})".format(", ".join(str(arg) for arg in op_meta.args))
 
     logger.info(
         "{0} {1} {2}".format(
@@ -236,7 +210,7 @@ def log_operation_start(op_meta: Dict, op_types: Optional[List] = None, prefix:
                 ),
                 "blue",
             ),
-            click.style(", ".join(op_meta["names"]), bold=True),
+            click.style(", ".join(op_meta.names), bold=True),
             args,
         ),
     )
@@ -244,7 +218,7 @@ def log_operation_start(op_meta: Dict, op_types: Optional[List] = None, prefix:
 
 def log_error_or_warning(
     host: "Host", ignore_errors: bool, description: str = "", continue_on_error: bool = False
-):
+) -> None:
     log_func = logger.error
     log_color = "red"
     log_text = "Error: " if description else "Error"
@@ -267,7 +241,7 @@ def log_error_or_warning(
     )
 
 
-def log_host_command_error(host: "Host", e, timeout: int = 0):
+def log_host_command_error(host: "Host", e: Exception, timeout: int = 0) -> None:
     if isinstance(e, timeout_error):
         logger.error(
             "{0}{1}".format(
@@ -327,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)
@@ -351,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 (
@@ -378,25 +364,20 @@ class get_file_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):
@@ -443,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}")