pyinfra 0.11.dev3__py3-none-any.whl → 3.5.1__py3-none-any.whl

pyinfra/api/inventory.py

@@ -1,18 +1,17 @@
-from collections import defaultdict
+from __future__ import annotations
 
-from pyinfra import logger
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Iterator
 
-from .connectors import (
-    ALL_CONNECTORS,
-    EXECUTION_CONNECTORS,
-    INVENTORY_CONNECTORS,
-)
+from .connectors import get_all_connectors, get_execution_connectors
 from .exceptions import NoConnectorError, NoGroupError, NoHostError
 from .host import Host
-from .util import FallbackDict
 
+if TYPE_CHECKING:
+    from pyinfra.api.state import State
 
-def extract_name_data(names):
+
+def extract_name_data(names: list[Any]):
     for name in names:
         data = {}
 
@@ -23,48 +22,32 @@ def extract_name_data(names):
         yield name, data
 
 
-class Inventory(object):
-    '''
+class Inventory:
+    """
     Represents a collection of target hosts. Stores and provides access to group data,
     host data and default data for these hosts.
 
     Args:
         names_data: tuple of ``(names, data)``
-        ssh_user: override SSH user
-        ssh_port: override SSH port
-        ssh_key: override SSH key filename
-        ssh_key_password: override password for the SSH key
-        ssh_password: override SSH password
-        **groups: map of group names -> ``(names, data)``
-    '''
-
-    state = None
-
-    def __init__(
-        self, names_data,
-        ssh_user=None, ssh_port=None, ssh_key=None,
-        ssh_key_password=None, ssh_password=None, **groups
-    ):
+        override_data: dictionary of data overrides
+        ssh_*: deprecated, use ``override_data.ssh_*``
+        winrm_*: deprecated, use ``override_data.winrm_*``
+        **groups: map of group name -> ``(names, data)``
+    """
+
+    state: "State"
+    groups: dict[str, list[Host]]
+
+    @staticmethod
+    def empty():
+        return 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
-
-        # In CLI mode these are --user, --key, etc
-        override_data = {
-            'ssh_user': ssh_user,
-            'ssh_key': ssh_key,
-            'ssh_key_password': ssh_key_password,
-            'ssh_port': ssh_port,
-            'ssh_password': ssh_password,
-        }
-        # Strip None values
-        override_data = {
-            key: value
-            for key, value in override_data.items()
-            if value is not None
-        }
-        self.override_data = override_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
 
@@ -72,11 +55,14 @@ class Inventory(object):
         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) -> None:
+        all_connectors = get_all_connectors()
+        execution_connectors = get_execution_connectors()
 
-    def make_hosts_and_groups(self, names, groups):
         # 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)
 
@@ -94,82 +80,61 @@ class Inventory(object):
         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"]
 
-            # Name is @connector?
-            if name[0] == '@':
+            if name[0] == "@":
                 connector_name = name[1:]
                 arg_string = None
 
-                if '/' in connector_name:
-                    connector_name, arg_string = connector_name.split('/', 1)
+                if "/" in connector_name:
+                    connector_name, arg_string = connector_name.split("/", 1)
 
-                if connector_name not in ALL_CONNECTORS:
+                if connector_name not in get_all_connectors():
                     raise NoConnectorError(
-                        'Invalid connector: {0}'.format(connector_name),
+                        "Invalid connector: {0}".format(connector_name),
                     )
 
-                # Execution connector? Simple, just set it for ther host
-                if connector_name in EXECUTION_CONNECTORS:
-                    executor = EXECUTION_CONNECTORS[connector_name]
-
-                # Inventory connector?
-                if connector_name in INVENTORY_CONNECTORS:
-                    logger.debug('Expanding inventory connector: {0}'.format(
-                        connector_name,
-                    ))
-
-                    for name, data, group_names in (
-                        INVENTORY_CONNECTORS[connector_name].make_names_data(arg_string)
-                    ):
-                        # Make a copy of the host data, update with any from
-                        # the connector.
-                        sub_host_data = host_data.copy()
-                        sub_host_data.update(data)
-
-                        # Assign the name/data/group_names from the connector
-                        self.host_data[name] = sub_host_data
-                        names_executors.append((name, executor))
-                        name_to_group_names[name].extend(group_names)
-
-                    continue
-
-            # Assign the name/data
-            self.host_data[name] = host_data
-            names_executors.append((name, executor))
+                # Execution connector? Simple, just set it for their host
+                if connector_name in execution_connectors:
+                    connector_cls = execution_connectors[connector_name]
+
+                names_data = all_connectors[connector_name].make_names_data(arg_string)
+                connector_inventory_name = name
+            else:
+                names_data = [(name, {}, [])]
+                connector_inventory_name = None
+
+            for sub_name, sub_data, sub_groups in names_data:
+                # Update any connector data with a copy of the host data (so that
+                # host data can override connector data).
+                sub_data.update(host_data.copy())
+
+                # Assign the name/data/groups from the connector
+                self.host_data[sub_name] = sub_data
+                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
+                # to the newly generated host name.
+                if connector_inventory_name:
+                    name_to_group_names[sub_name].extend(
+                        name_to_group_names[connector_inventory_name],
+                    )
 
         # Now we can actually make Host instances
-        hosts = {}
+        hosts: dict[str, "Host"] = {}
 
-        for name, executor in names_executors:
+        for name, connector_cls in names_connectors:
             host_groups = name_to_group_names[name]
 
-            # Create the (waterfall data: override, host, group, global)
-            host_data = FallbackDict(
-                self.get_override_data(),
-                self.get_host_data(name),
-                self.get_groups_data(host_groups),
-                self.get_data(),
-                # Pass the method, rather than data, as this comes from the
-                # state and can change during deploy(s).
-                self.get_deploy_data,
-            )
-
-            # Create the Host object
-            host = Host(
-                name,
-                inventory=self,
-                groups=name_to_group_names.get(name),
-                data=host_data,
-                executor=executor,
-            )
+            host = Host(name, inventory=self, groups=host_groups, connector_cls=connector_cls)
             hosts[name] = host
 
             # And push into any groups
@@ -177,101 +142,107 @@ class Inventory(object):
                 if host not in self.groups[group_name]:
                     self.groups[group_name].append(host)
 
-        return hosts
-
-    def __len__(self):
-        '''
-        Returns the number of active inventory hosts.
-        '''
+        self.hosts = hosts
 
-        if not self.state or not self.state.active_hosts:
-            return len(self.hosts)
+    def __len__(self) -> int:
+        """
+        Returns the number of inventory hosts.
+        """
 
-        return len(self.state.active_hosts)
+        return len(self.hosts)
 
-    def __iter__(self):
-        '''
-        Iterates over active inventory hosts.
-        '''
+    def __iter__(self) -> Iterator["Host"]:
+        """
+        Iterates over all inventory hosts.
+        """
 
-        if not self.state or not self.state.active_hosts:
-            return self.iter_all_hosts()
+        return iter(self.hosts.values())
 
+    def iter_active_hosts(self) -> Iterator["Host"]:
+        """
+        Iterates over active inventory hosts.
+        """
         return iter(self.state.active_hosts)
 
-    def len_all_hosts(self):
-        '''
-        Returns the number of hosts in the inventory, active or not.
-        '''
-
-        return len(self.hosts)
+    def len_active_hosts(self) -> int:
+        """
+        Returns the number of active inventory hosts.
+        """
+        return len(self.state.active_hosts)
 
-    def iter_all_hosts(self):
-        '''
-        Iterates over all inventory hosts, active or not.
-        '''
+    def iter_activated_hosts(self) -> Iterator["Host"]:
+        """
+        Iterates over activated inventory hosts.
+        """
+        return iter(self.state.activated_hosts)
 
-        return iter(self.hosts.values())
+    def len_activated_hosts(self) -> int:
+        """
+        Returns the number of activated inventory hosts.
+        """
+        return len(self.state.activated_hosts)
 
-    def get_host(self, name, default=NoHostError):
-        '''
+    def get_host(self, name: str, default=NoHostError) -> Host:
+        """
         Get a single host by name.
-        '''
+        """
 
         if name in self.hosts:
             return self.hosts[name]
 
         if default is NoHostError:
-            raise NoHostError('No such host: {0}'.format(name))
+            raise NoHostError("No such host: {0}".format(name))
 
+        # TODO: remove default here?
         return default
 
-    def get_group(self, name, default=NoGroupError):
-        '''
+    def get_group(self, name: str, default=NoGroupError) -> list[Host]:
+        """
         Get a list of hosts belonging to a group.
-        '''
+        """
 
         if name in self.groups:
             return self.groups[name]
 
         if default is NoGroupError:
-            raise NoGroupError('No such group: {0}'.format(name))
+            raise NoGroupError("No such group: {0}".format(name))
 
+        # TODO: remove default here?
         return default
 
     def get_data(self):
-        '''
+        """
         Get the base/all data attached to this inventory.
-        '''
+        """
 
         return self.data
 
     def get_override_data(self):
-        '''
+        """
         Get override data for this inventory.
-        '''
+        """
 
         return self.override_data
 
-    def get_host_data(self, hostname):
-        '''
+    def get_host_data(self, hostname: str):
+        """
         Get data for a single host in this inventory.
-        '''
+        """
 
         return self.host_data.get(hostname, {})
 
     def get_group_data(self, group):
-        '''
+        """
         Get data for a single group in this inventory.
-        '''
+        """
 
         return self.group_data.get(group, {})
 
     def get_groups_data(self, groups):
-        '''
+        """
         Gets aggregated data from a list of groups. Vars are collected in order so, for
         any groups which define the same var twice, the last group's value will hold.
-        '''
+        """
 
         data = {}
 
@@ -279,13 +250,3 @@ class Inventory(object):
             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 {}