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

pyinfra/api/host.py

@@ -1,22 +1,39 @@
+from __future__ import annotations
+
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generator,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+from uuid import uuid4
 
 import click
-from gevent.lock import BoundedSemaphore
+from typing_extensions import Unpack
 
 from pyinfra import logger
-from pyinfra.connectors.util import remove_any_sudo_askpass_file
+from pyinfra.connectors.base import BaseConnector
+from pyinfra.connectors.util import CommandOutput, remove_any_sudo_askpass_file
 
 from .connectors import get_execution_connector
 from .exceptions import ConnectError
-from .facts import create_host_fact, delete_host_fact, get_host_fact, reload_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
     from pyinfra.api.inventory import Inventory
     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.
@@ -31,7 +48,7 @@ class HostData:
     Combines multiple AttrData's to search for attributes.
     """
 
-    override_datas: Dict[str, Any]
+    override_datas: dict[str, Any]
 
     def __init__(self, host: "Host", *datas):
         self.__dict__["host"] = host
@@ -82,28 +99,36 @@ class Host:
     data.
     """
 
-    connection = None
     state: "State"
+    connector_cls: type[BaseConnector]
+    connector: BaseConnector
+    connected: bool = False
 
     # Current context inside an @operation function (op gen stage)
     in_op: bool = False
+    in_callback_op: bool = False
     current_op_hash: Optional[str] = None
-    current_op_global_kwargs: Dict[str, Any]
+    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]
+    loop_position: list[int]
 
     # Arbitrary data dictionary connectors may use
-    connector_data: Dict[str, Any]
+    connector_data: dict[str, Any]
 
     def loop(self, iterable):
         self.loop_position.append(0)
@@ -117,26 +142,20 @@ class Host:
         name: str,
         inventory: "Inventory",
         groups,
-        executor=get_execution_connector("ssh"),
+        connector_cls=get_execution_connector("ssh"),
     ):
         self.inventory = inventory
         self.groups = groups
-        self.executor = executor
+        self.connector_cls = connector_cls
         self.name = name
 
         self.loop_position = []
 
         self.connector_data = {}
-        self.current_op_global_kwargs = {}
-
-        # Fact data store
-        # TODO: how to not have Any here?
-        self.facts: Dict[str, Any] = {}
-        self.facts_lock = BoundedSemaphore()
 
         # Append only list of operation hashes as called on this host, used to
         # generate a DAG to create the final operation order.
-        self.op_hash_order: List[str] = []
+        self.op_hash_order: list[str] = []
 
         # Create the (waterfall data: override, host, group, global)
         self.data = HostData(
@@ -149,16 +168,20 @@ class Host:
             self.get_deploy_data,
         )
 
+    def init(self, state: "State") -> None:
+        self.state = state
+        self.connector = self.connector_cls(state, self)
+
+        longest_name_len = max([len(host.name) for host in self.inventory])
+        padding_diff = longest_name_len - len(self.name)
+        self.print_prefix_padding = "".join(" " for _ in range(0, padding_diff))
+
     def __str__(self):
         return "{0}".format(self.name)
 
     def __repr__(self):
         return "Host({0})".format(self.name)
 
-    @property
-    def connected(self) -> bool:
-        return self.connection is not None
-
     @property
     def host_data(self):
         return self.inventory.get_host_data(self.name)
@@ -168,30 +191,37 @@ class Host:
         return self.inventory.get_groups_data(self.groups)
 
     @property
-    def print_prefix(self):
+    def print_prefix(self) -> str:
         if self.nested_executing_op_hash:
-            return "{0}[{1}] {2} ".format(
+            return "{0}[{1}] {2}{3} ".format(
                 click.style(""),  # reset
                 click.style(self.name, bold=True),
                 click.style("nested", "blue"),
+                self.print_prefix_padding,
             )
 
-        return "{0}[{1}] ".format(
+        return "{0}[{1}]{2} ".format(
             click.style(""),  # reset
             click.style(self.name, bold=True),
+            self.print_prefix_padding,
         )
 
-    def style_print_prefix(self, *args, **kwargs):
-        return "{0}[{1}] ".format(
+    def style_print_prefix(self, *args, **kwargs) -> str:
+        return "{0}[{1}]{2} ".format(
             click.style(""),  # reset
             click.style(self.name, *args, **kwargs),
+            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}")
+
+    def log_styled(self, message, log_func=logger.info, **kwargs):
+        message_styled = click.style(message, **kwargs)
+        self.log(message_styled, log_func=log_func)
 
-        return {}
+    def get_deploy_data(self):
+        return self.current_op_deploy_data or self.current_deploy_data or {}
 
     def noop(self, description):
         """
@@ -201,8 +231,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``.
@@ -219,6 +266,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
@@ -245,33 +299,55 @@ class Host:
             old_deploy_data,
         )
 
-    # Host facts
-    #
+    @memoize
+    def _get_temp_directory(self):
+        temp_directory = self.state.config.TEMP_DIR
 
-    def get_fact(self, name_or_cls, *args, **kwargs):
-        """
-        Get a fact for this host, reading from the cache if present.
-        """
-        return get_host_fact(self.state, self, name_or_cls, args=args, kwargs=kwargs)
+        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
 
-    def reload_fact(self, name_or_cls, *args, **kwargs):
-        """
-        Get a fact for this host without using any cached value, always re-fetch the fact data
-        from the host and then cache it.
-        """
-        return reload_host_fact(self.state, self, name_or_cls, args=args, kwargs=kwargs)
+            temp_directory = self.get_fact(TmpDir)
 
-    def create_fact(self, name_or_cls, data=None, kwargs=None):
+        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):
         """
-        Create a new fact for this host in the fact cache.
+        Generate a temporary filename for this deploy.
         """
-        return create_host_fact(self.state, self, name_or_cls, data, kwargs=kwargs)
 
-    def delete_fact(self, name_or_cls, kwargs=None):
+        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):
         """
-        Remove an existing fact for this host in the fact cache.
+        Get a fact for this host, reading from the cache if present.
         """
-        return delete_host_fact(self.state, self, name_or_cls, kwargs=kwargs)
+        return get_host_fact(self.state, self, name_or_cls, args=args, kwargs=kwargs)
 
     # Connector proxy
     #
@@ -286,11 +362,11 @@ class Host:
         """
 
         self._check_state()
-        if not self.connection:
+        if not self.connected:
             self.state.trigger_callbacks("host_before_connect", self)
 
             try:
-                self.connection = self.executor.connect(self.state, self)
+                self.connector.connect()
             except ConnectError as e:
                 if show_errors:
                     log_message = "{0}{1}".format(
@@ -316,8 +392,7 @@ class Host:
 
                 logger.info(log_message)
                 self.state.trigger_callbacks("host_connect", self)
-
-        return self.connection
+                self.connected = True
 
     def disconnect(self):
         """
@@ -325,50 +400,43 @@ class Host:
         """
         self._check_state()
 
-        # Disconnect is an optional function for executors if needed
-        disconnect_func = getattr(self.executor, "disconnect", None)
+        # Disconnect is an optional function for connectors if needed
+        disconnect_func = getattr(self.connector, "disconnect", None)
         if disconnect_func:
-            return disconnect_func(self.state, self)
+            return disconnect_func()
 
         # TODO: consider whether this should be here!
         remove_any_sudo_askpass_file(self)
 
         self.state.trigger_callbacks("host_disconnect", self)
 
-    def run_shell_command(self, *args, **kwargs):
+    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.
         """
         self._check_state()
-        return self.executor.run_shell_command(self.state, self, *args, **kwargs)
+        return self.connector.run_shell_command(*args, **kwargs)
 
-    def put_file(self, *args, **kwargs):
+    def put_file(self, *args, **kwargs) -> bool:
         """
         Low level method to upload a file to the host via it's configured connector.
         """
         self._check_state()
-        return self.executor.put_file(self.state, self, *args, **kwargs)
+        return self.connector.put_file(*args, **kwargs)
 
-    def get_file(self, *args, **kwargs):
+    def get_file(self, *args, **kwargs) -> bool:
         """
         Low level method to download a file from the host via it's configured connector.
         """
         self._check_state()
-        return self.executor.get_file(self.state, self, *args, **kwargs)
-
-    # Rsync - optional executor specific ability
+        return self.connector.get_file(*args, **kwargs)
 
-    def check_can_rsync(self):
-        check_can_rsync_func = getattr(self.executor, "check_can_rsync", None)
-        if check_can_rsync_func:
-            return check_can_rsync_func(self)
+    # Rsync - optional connector specific ability
 
-        raise NotImplementedError(
-            "The {0} connector does not support rsync!".format(
-                self.executor.__name__,
-            ),
-        )
+    def check_can_rsync(self) -> None:
+        self._check_state()
+        return self.connector.check_can_rsync()
 
-    def rsync(self, *args, **kwargs):
+    def rsync(self, *args, **kwargs) -> bool:
         self._check_state()
-        return self.executor.rsync(self.state, self, *args, **kwargs)
+        return self.connector.rsync(*args, **kwargs)

pyinfra/api/inventory.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from collections import defaultdict
-from typing import TYPE_CHECKING, Any, Iterator, List
+from typing import TYPE_CHECKING, Any, Iterator
 
 from .connectors import get_all_connectors, get_execution_connectors
 from .exceptions import NoConnectorError, NoGroupError, NoHostError
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
     from pyinfra.api.state import State
 
 
-def extract_name_data(names: List[Any]):
+def extract_name_data(names: list[Any]):
     for name in names:
         data = {}
 
@@ -38,8 +40,8 @@ class Inventory:
     def __init__(self, names_data, override_data=None, **groups):
         # Setup basics
         self.groups = defaultdict(list)  # lists of Host objects
-        self.host_data = defaultdict(dict)  # dict of name -> data
-        self.group_data = defaultdict(dict)  # dict of name -> data
+        self.host_data: dict[str, dict] = defaultdict(dict)  # dict of name -> data
+        self.group_data: dict[str, dict] = defaultdict(dict)  # dict of name -> data
         self.override_data = override_data or {}
 
         names, data = names_data
@@ -48,14 +50,14 @@ class Inventory:
         self.data = data
 
         # Create the actual host instances and groups
-        self.hosts = self.make_hosts_and_groups(names, groups)
+        self.make_hosts_and_groups(names, groups)
 
-    def make_hosts_and_groups(self, names, groups):
+    def make_hosts_and_groups(self, names, groups) -> None:
         all_connectors = get_all_connectors()
         execution_connectors = get_execution_connectors()
 
         # Map name -> data
-        name_to_data = defaultdict(dict)
+        name_to_data: dict[str, dict] = defaultdict(dict)
         # Map name -> group names
         name_to_group_names = defaultdict(list)
 
@@ -73,14 +75,14 @@ class Inventory:
         for name, data in extract_name_data(names):
             name_to_data[name].update(data)
 
-        # Now, use the above to fill self.host_data and populate names_executors
-        names_executors = []
+        # Now, use the above to fill self.host_data and populate names_connectors
+        names_connectors = []
 
         for name, _ in extract_name_data(names):
             host_data = name_to_data[name]
 
             # Default to executing commands with the ssh connector
-            executor = execution_connectors["ssh"]
+            connector_cls = execution_connectors["ssh"]
 
             if name[0] == "@":
                 connector_name = name[1:]
@@ -96,7 +98,7 @@ class Inventory:
 
                 # Execution connector? Simple, just set it for their host
                 if connector_name in execution_connectors:
-                    executor = execution_connectors[connector_name]
+                    connector_cls = execution_connectors[connector_name]
 
                 names_data = all_connectors[connector_name].make_names_data(arg_string)
                 connector_inventory_name = name
@@ -111,7 +113,7 @@ class Inventory:
 
                 # Assign the name/data/groups from the connector
                 self.host_data[sub_name] = sub_data
-                names_executors.append((sub_name, executor))
+                names_connectors.append((sub_name, connector_cls))
                 name_to_group_names[sub_name].extend(sub_groups)
 
                 # If we have a connector inventory name, copy any groups attached
@@ -124,10 +126,10 @@ class Inventory:
         # Now we can actually make Host instances
         hosts: dict[str, "Host"] = {}
 
-        for name, executor in names_executors:
+        for name, connector_cls in names_connectors:
             host_groups = name_to_group_names[name]
 
-            host = Host(name, inventory=self, groups=host_groups, executor=executor)
+            host = Host(name, inventory=self, groups=host_groups, connector_cls=connector_cls)
             hosts[name] = host
 
             # And push into any groups
@@ -135,7 +137,7 @@ class Inventory:
                 if host not in self.groups[group_name]:
                     self.groups[group_name].append(host)
 
-        return hosts
+        self.hosts = hosts
 
     def __len__(self) -> int:
         """
@@ -241,13 +243,3 @@ class Inventory:
             data.update(self.get_group_data(group))
 
         return data
-
-    def get_deploy_data(self):
-        """
-        Gets any default data attached to the current deploy, if any.
-        """
-
-        if self.state and self.state.deploy_data:
-            return self.state.deploy_data
-
-        return {}