pyinfra 3.0b0__py2.py3-none-any.whl → 3.0b1__py2.py3-none-any.whl

pyinfra/connectors/local.py

@@ -25,7 +25,8 @@ if TYPE_CHECKING:
 
 class LocalConnector(BaseConnector):
     """
-    The ``@local`` connector executes changes on the local machine using subprocesses. **This connector is only compatible with MacOS & Linux hosts**.
+    The ``@local`` connector executes changes on the local machine using
+    subprocesses. **This connector is only compatible with MacOS & Linux hosts**.
 
     Examples:
 
@@ -38,8 +39,8 @@ class LocalConnector(BaseConnector):
     handles_execution = True
 
     @staticmethod
-    def make_names_data(_=None):
-        if _ is not None:
+    def make_names_data(name=None):
+        if name is not None:
             raise InventoryError("Cannot have more than one @local")
 
         yield "@local", {}, ["@local"]
@@ -205,8 +206,7 @@ class LocalConnector(BaseConnector):
 
         return True
 
-    @staticmethod
-    def check_can_rsync(host):
+    def check_can_rsync(self):
         if not find_executable("rsync"):
             raise NotImplementedError("The `rsync` binary is not available on this system.")
 

pyinfra/connectors/ssh.py

@@ -2,7 +2,9 @@ from __future__ import annotations
 
 import shlex
 from distutils.spawn import find_executable
+from random import uniform
 from socket import error as socket_error, gaierror
+from time import sleep
 from typing import TYPE_CHECKING, Any, Iterable, Optional, Tuple
 
 import click
@@ -48,6 +50,10 @@ class ConnectorData(TypedDict):
 
     ssh_paramiko_connect_kwargs: dict
 
+    ssh_connect_retries: int
+    ssh_connect_retry_min_delay: float
+    ssh_connect_retry_max_delay: float
+
 
 connector_data_meta: dict[str, DataMeta] = {
     "ssh_hostname": DataMeta("SSH hostname"),
@@ -77,6 +83,15 @@ connector_data_meta: dict[str, DataMeta] = {
     "ssh_paramiko_connect_kwargs": DataMeta(
         "Override keyword arguments passed into Paramiko's ``SSHClient.connect``"
     ),
+    "ssh_connect_retries": DataMeta("Number of tries to connect via ssh", 0),
+    "ssh_connect_retry_min_delay": DataMeta(
+        "Lower bound for random delay between retries",
+        0.1,
+    ),
+    "ssh_connect_retry_max_delay": DataMeta(
+        "Upper bound for random delay between retries",
+        0.5,
+    ),
 }
 
 
@@ -125,8 +140,9 @@ class SSHConnector(BaseConnector):
 
     client: Optional[SSHClient] = None
 
-    def make_names_data(hostname):
-        yield "@ssh/{0}".format(hostname), {"ssh_hostname": hostname}, []
+    @staticmethod
+    def make_names_data(name):
+        yield "@ssh/{0}".format(name), {"ssh_hostname": name}, []
 
     def make_paramiko_kwargs(self) -> dict[str, Any]:
         kwargs = {
@@ -172,6 +188,29 @@ class SSHConnector(BaseConnector):
         return kwargs
 
     def connect(self) -> None:
+        retries = self.data["ssh_connect_retries"]
+
+        try:
+            while True:
+                try:
+                    return self._connect()
+                except (SSHException, gaierror, socket_error, EOFError):
+                    if retries == 0:
+                        raise
+                    retries -= 1
+                    min_delay = self.data["ssh_connect_retry_min_delay"]
+                    max_delay = self.data["ssh_connect_retry_max_delay"]
+                    sleep(uniform(min_delay, max_delay))
+        except SSHException as e:
+            raise_connect_error(self.host, "SSH error", e)
+        except gaierror as e:
+            raise_connect_error(self.host, "Could not resolve hostname", e)
+        except socket_error as e:
+            raise_connect_error(self.host, "Could not connect", e)
+        except EOFError as e:
+            raise_connect_error(self.host, "EOF error", e)
+
+    def _connect(self) -> None:
         """
         Connect to a single host. Returns the SSH client if successful. Stateless by
         design so can be run in parallel.
@@ -221,18 +260,6 @@ class SSHConnector(BaseConnector):
                 f"Host key for {e.hostname} does not match.",
             )
 
-        except SSHException as e:
-            raise_connect_error(self.host, "SSH error", e)
-
-        except gaierror:
-            raise_connect_error(self.host, "Could not resolve hostname", hostname)
-
-        except socket_error as e:
-            raise_connect_error(self.host, "Could not connect", e)
-
-        except EOFError as e:
-            raise_connect_error(self.host, "EOF error", e)
-
     def run_shell_command(
         self,
         command: StringCommand,
@@ -450,16 +477,10 @@ class SSHConnector(BaseConnector):
             self._put_file(filename_or_io, temp_file)
 
             # Make sure our sudo/su user can access the file
-            if _su_user:
-                command = StringCommand("setfacl", "-m", "u:{0}:r".format(_su_user), temp_file)
-            elif _sudo_user:
-                command = StringCommand("setfacl", "-m", "u:{0}:r".format(_sudo_user), temp_file)
-            elif _doas_user:
-                command = StringCommand("setfacl", "-m", "u:{0}:r".format(_doas_user), temp_file)
-
-            if _su_user or _sudo_user or _doas_user:
+            other_user = _su_user or _sudo_user or _doas_user
+            if other_user:
                 status, output = self.run_shell_command(
-                    command,
+                    StringCommand("setfacl", "-m", f"u:{other_user}:r", temp_file),
                     print_output=print_output,
                     print_input=print_input,
                     **arguments,

pyinfra/connectors/terraform.py

@@ -28,7 +28,8 @@ def _flatten_dict(d: dict, parent_key: str = "", sep: str = "."):
 
 class TerraformInventoryConnector(BaseConnector):
     """
-    Generate one or more SSH hosts from a Terraform output variable. The variable must be a list of hostnames or dictionaries.
+    Generate one or more SSH hosts from a Terraform output variable. The variable
+    must be a list of hostnames or dictionaries.
 
     Output is fetched from a flattened JSON dictionary output from ``terraform output
     -json``. For example the following object:
@@ -77,21 +78,23 @@ class TerraformInventoryConnector(BaseConnector):
     """
 
     @staticmethod
-    def make_names_data(output_key=None):
+    def make_names_data(name=None):
         show_warning()
 
-        if not output_key:
-            raise InventoryError("No Terraform output key!")
+        if not name:
+            name = ""
 
         with progress_spinner({"fetch terraform output"}):
             tf_output_raw = local.shell("terraform output -json")
 
+        assert isinstance(tf_output_raw, str)
         tf_output = json.loads(tf_output_raw)
         tf_output = _flatten_dict(tf_output)
 
-        tf_output_value = tf_output.get(output_key)
+        tf_output_value = tf_output.get(name)
         if tf_output_value is None:
-            raise InventoryError(f"No Terraform output with key: `{output_key}`")
+            keys = "\n".join(f"   - {k}" for k in tf_output.keys())
+            raise InventoryError(f"No Terraform output with key: `{name}`, valid keys:\n{keys}")
 
         if not isinstance(tf_output_value, list):
             raise InventoryError(

pyinfra/connectors/util.py

@@ -40,7 +40,7 @@ def run_local_process(
     stdin=None,
     timeout: Optional[int] = None,
     print_output: bool = False,
-    print_prefix=None,
+    print_prefix: str = "",
 ) -> tuple[int, "CommandOutput"]:
     process = Popen(command, shell=True, stdout=PIPE, stderr=PIPE, stdin=PIPE)
 

pyinfra/connectors/vagrant.py

@@ -132,8 +132,8 @@ class VagrantInventoryConnector(BaseConnector):
     """
 
     @staticmethod
-    def make_names_data(limit=None):
-        vagrant_ssh_info = get_vagrant_config(limit)
+    def make_names_data(name=None):
+        vagrant_ssh_info = get_vagrant_config(name)
 
         logger.debug("Got Vagrant SSH info: \n%s", vagrant_ssh_info)
 
@@ -170,10 +170,11 @@ class VagrantInventoryConnector(BaseConnector):
             hosts.append(_make_name_data(current_host))
 
         if not hosts:
-            if limit:
+            if name:
                 raise InventoryError(
-                    "No running Vagrant instances matching `{0}` found!".format(limit)
+                    "No running Vagrant instances matching `{0}` found!".format(name)
                 )
             raise InventoryError("No running Vagrant instances found!")
 
-        return hosts
+        for host in hosts:
+            yield host

pyinfra/facts/choco.py

@@ -16,7 +16,7 @@ class ChocoPackages(FactBase):
         }
     """
 
-    command = "choco list --local-only"
+    command = "choco list"
     shell_executable = "ps"
 
     default = dict

pyinfra/facts/deb.py

@@ -48,8 +48,8 @@ class DebPackage(FactBase):
     """
 
     _regexes = {
-        "name": r"^Package: ({0})$".format(DEB_PACKAGE_NAME_REGEX),
-        "version": r"^Version: ({0})$".format(DEB_PACKAGE_VERSION_REGEX),
+        "name": r"^Package:\s+({0})$".format(DEB_PACKAGE_NAME_REGEX),
+        "version": r"^Version:\s+({0})$".format(DEB_PACKAGE_VERSION_REGEX),
     }
 
     requires_command = "dpkg"

pyinfra/facts/postgres.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+from pyinfra.api import FactBase, MaskString, QuoteString, StringCommand
+from pyinfra.api.util import try_int
+
+from .util.databases import parse_columns_and_rows
+
+
+def make_psql_command(
+    database=None,
+    user=None,
+    password=None,
+    host=None,
+    port=None,
+    executable="psql",
+):
+    target_bits: list[str] = []
+
+    if password:
+        target_bits.append(MaskString('PGPASSWORD="{0}"'.format(password)))
+
+    target_bits.append(executable)
+
+    if database:
+        target_bits.append("-d {0}".format(database))
+
+    if user:
+        target_bits.append("-U {0}".format(user))
+
+    if host:
+        target_bits.append("-h {0}".format(host))
+
+    if port:
+        target_bits.append("-p {0}".format(port))
+
+    return StringCommand(*target_bits)
+
+
+def make_execute_psql_command(command, **psql_kwargs):
+    return StringCommand(
+        make_psql_command(**psql_kwargs),
+        "-Ac",
+        QuoteString(command),  # quote this whole item as a single shell argument
+    )
+
+
+class PostgresFactBase(FactBase):
+    abstract = True
+
+    psql_command: str
+    requires_command = "psql"
+
+    def command(
+        self,
+        psql_user=None,
+        psql_password=None,
+        psql_host=None,
+        psql_port=None,
+    ):
+        return make_execute_psql_command(
+            self.psql_command,
+            user=psql_user,
+            password=psql_password,
+            host=psql_host,
+            port=psql_port,
+        )
+
+
+class PostgresRoles(PostgresFactBase):
+    """
+    Returns a dict of PostgreSQL roles and data:
+
+    .. code:: python
+
+        {
+            "pyinfra": {
+                "super": true,
+                "createrole": false,
+                "createdb": false,
+                ...
+            },
+        }
+    """
+
+    default = dict
+    psql_command = "SELECT * FROM pg_catalog.pg_roles"
+
+    def process(self, output):
+        # Remove the last line of the output (row count)
+        output = output[:-1]
+        rows = parse_columns_and_rows(
+            output,
+            "|",
+            # Remove the "rol" prefix on column names
+            remove_column_prefix="rol",
+        )
+
+        users = {}
+
+        for details in rows:
+            for key, value in list(details.items()):
+                if key in ("oid", "connlimit"):
+                    details[key] = try_int(value)
+
+                if key in (
+                    "super",
+                    "inherit",
+                    "createrole",
+                    "createdb",
+                    "canlogin",
+                    "replication",
+                    "bypassrls",
+                ):
+                    details[key] = value == "t"
+
+            users[details.pop("name")] = details
+
+        return users
+
+
+class PostgresDatabases(PostgresFactBase):
+    """
+    Returns a dict of PostgreSQL databases and metadata:
+
+    .. code:: python
+
+        {
+            "pyinfra_stuff": {
+                "encoding": "UTF8",
+                "collate": "en_US.UTF-8",
+                "ctype": "en_US.UTF-8",
+                ...
+            },
+        }
+    """
+
+    default = dict
+    psql_command = "SELECT pg_catalog.pg_encoding_to_char(encoding), * FROM pg_catalog.pg_database"
+
+    def process(self, output):
+        # Remove the last line of the output (row count)
+        output = output[:-1]
+        rows = parse_columns_and_rows(
+            output,
+            "|",
+            # Remove the "dat" prefix on column names
+            remove_column_prefix="dat",
+        )
+
+        databases = {}
+
+        for details in rows:
+            details["encoding"] = details.pop("pg_encoding_to_char")
+
+            for key, value in list(details.items()):
+                if key.endswith("id") or key in (
+                    "dba",
+                    "tablespace",
+                    "connlimit",
+                ):
+                    details[key] = try_int(value)
+
+                if key in ("istemplate", "allowconn"):
+                    details[key] = value == "t"
+
+            databases[details.pop("name")] = details
+
+        return databases

pyinfra/facts/postgresql.py

@@ -1,168 +1,9 @@
-from __future__ import annotations
+from .postgres import PostgresDatabases, PostgresRoles
 
-from pyinfra.api import FactBase, MaskString, QuoteString, StringCommand
-from pyinfra.api.util import try_int
 
-from .util.databases import parse_columns_and_rows
+class PostgresqlRoles(PostgresRoles):
+    deprecated = True
 
 
-def make_psql_command(
-    database=None,
-    user=None,
-    password=None,
-    host=None,
-    port=None,
-    executable="psql",
-):
-    target_bits: list[str] = []
-
-    if password:
-        target_bits.append(MaskString('PGPASSWORD="{0}"'.format(password)))
-
-    target_bits.append(executable)
-
-    if database:
-        target_bits.append("-d {0}".format(database))
-
-    if user:
-        target_bits.append("-U {0}".format(user))
-
-    if host:
-        target_bits.append("-h {0}".format(host))
-
-    if port:
-        target_bits.append("-p {0}".format(port))
-
-    return StringCommand(*target_bits)
-
-
-def make_execute_psql_command(command, **psql_kwargs):
-    return StringCommand(
-        make_psql_command(**psql_kwargs),
-        "-Ac",
-        QuoteString(command),  # quote this whole item as a single shell argument
-    )
-
-
-class PostgresqlFactBase(FactBase):
-    abstract = True
-
-    psql_command: str
-    requires_command = "psql"
-
-    def command(
-        self,
-        psql_user=None,
-        psql_password=None,
-        psql_host=None,
-        psql_port=None,
-    ):
-        return make_execute_psql_command(
-            self.psql_command,
-            user=psql_user,
-            password=psql_password,
-            host=psql_host,
-            port=psql_port,
-        )
-
-
-class PostgresqlRoles(PostgresqlFactBase):
-    """
-    Returns a dict of PostgreSQL roles and data:
-
-    .. code:: python
-
-        {
-            "pyinfra": {
-                "super": true,
-                "createrole": false,
-                "createdb": false,
-                ...
-            },
-        }
-    """
-
-    default = dict
-    psql_command = "SELECT * FROM pg_catalog.pg_roles"
-
-    def process(self, output):
-        # Remove the last line of the output (row count)
-        output = output[:-1]
-        rows = parse_columns_and_rows(
-            output,
-            "|",
-            # Remove the "rol" prefix on column names
-            remove_column_prefix="rol",
-        )
-
-        users = {}
-
-        for details in rows:
-            for key, value in list(details.items()):
-                if key in ("oid", "connlimit"):
-                    details[key] = try_int(value)
-
-                if key in (
-                    "super",
-                    "inherit",
-                    "createrole",
-                    "createdb",
-                    "canlogin",
-                    "replication",
-                    "bypassrls",
-                ):
-                    details[key] = value == "t"
-
-            users[details.pop("name")] = details
-
-        return users
-
-
-class PostgresqlDatabases(PostgresqlFactBase):
-    """
-    Returns a dict of PostgreSQL databases and metadata:
-
-    .. code:: python
-
-        {
-            "pyinfra_stuff": {
-                "encoding": "UTF8",
-                "collate": "en_US.UTF-8",
-                "ctype": "en_US.UTF-8",
-                ...
-            },
-        }
-    """
-
-    default = dict
-    psql_command = "SELECT pg_catalog.pg_encoding_to_char(encoding), * FROM pg_catalog.pg_database"
-
-    def process(self, output):
-        # Remove the last line of the output (row count)
-        output = output[:-1]
-        rows = parse_columns_and_rows(
-            output,
-            "|",
-            # Remove the "dat" prefix on column names
-            remove_column_prefix="dat",
-        )
-
-        databases = {}
-
-        for details in rows:
-            details["encoding"] = details.pop("pg_encoding_to_char")
-
-            for key, value in list(details.items()):
-                if key.endswith("id") or key in (
-                    "dba",
-                    "tablespace",
-                    "connlimit",
-                ):
-                    details[key] = try_int(value)
-
-                if key in ("istemplate", "allowconn"):
-                    details[key] = value == "t"
-
-            databases[details.pop("name")] = details
-
-        return databases
+class PostgresqlDatabases(PostgresDatabases):
+    deprecated = True

pyinfra/facts/systemd.py

@@ -1,7 +1,7 @@
 import re
-from typing import Dict
+from typing import Dict, Iterable
 
-from pyinfra.api import FactBase
+from pyinfra.api import FactBase, FactTypeError, QuoteString, StringCommand
 
 # Valid unit names consist of a "name prefix" and a dot and a suffix specifying the unit type.
 # The "unit prefix" must consist of one or more valid characters
@@ -22,18 +22,16 @@ SYSTEMD_UNIT_NAME_REGEX = (
 
 def _make_systemctl_cmd(user_mode=False, machine=None, user_name=None):
     # base command for normal and user mode
-    systemctl_cmd = "systemctl --user" if user_mode else "systemctl"
+    systemctl_cmd = ["systemctl --user"] if user_mode else ["systemctl"]
 
     # add user and machine flag if given in args
     if machine is not None:
         if user_name is not None:
-            machine_opt = "--machine={1}@{0}".format(machine, user_name)
+            systemctl_cmd.append("--machine={1}@{0}".format(machine, user_name))
         else:
-            machine_opt = "--machine={0}".format(machine)
+            systemctl_cmd.append("--machine={0}".format(machine))
 
-        systemctl_cmd = "{0} {1}".format(systemctl_cmd, machine_opt)
-
-    return systemctl_cmd
+    return StringCommand(*systemctl_cmd)
 
 
 class SystemdStatus(FactBase[Dict[str, bool]]):
@@ -59,14 +57,32 @@ class SystemdStatus(FactBase[Dict[str, bool]]):
     state_key = "SubState"
     state_values = ["running", "waiting", "exited"]
 
-    def command(self, user_mode=False, machine=None, user_name=None):
+    def command(self, user_mode=False, machine=None, user_name=None, services=None):
         fact_cmd = _make_systemctl_cmd(
             user_mode=user_mode,
             machine=machine,
             user_name=user_name,
         )
 
-        return f"{fact_cmd} show --all --property Id --property {self.state_key} '*'"
+        if services is None:
+            service_strs = [QuoteString("*")]
+        elif isinstance(services, str):
+            service_strs = [QuoteString(services)]
+        elif isinstance(services, Iterable):
+            service_strs = [QuoteString(s) for s in services]
+        else:
+            raise FactTypeError(f"Invalid type passed for services argument: {type(services)}")
+
+        return StringCommand(
+            fact_cmd,
+            "show",
+            "--all",
+            "--property",
+            "Id",
+            "--property",
+            self.state_key,
+            *service_strs,
+        )
 
     def process(self, output) -> Dict[str, bool]:
         services: Dict[str, bool] = {}

pyinfra/operations/files.py

@@ -11,6 +11,7 @@ import traceback
 from datetime import timedelta
 from fnmatch import fnmatch
 from io import StringIO
+from pathlib import Path
 from typing import Union
 
 from jinja2 import TemplateRuntimeError, TemplateSyntaxError, UndefinedError
@@ -30,6 +31,7 @@ from pyinfra.api import (
 from pyinfra.api.command import make_formatted_string_command
 from pyinfra.api.util import (
     get_call_location,
+    get_file_io,
     get_file_sha1,
     get_path_permissions_mode,
     get_template,
@@ -569,7 +571,7 @@ def sync(
     put_files = []
     ensure_dirnames = []
     for dirpath, dirnames, filenames in os.walk(src, topdown=True):
-        remote_dirpath = os.path.normpath(os.path.relpath(dirpath, src))
+        remote_dirpath = Path(os.path.normpath(os.path.relpath(dirpath, src))).as_posix()
 
         # Filter excluded dirs
         for child_dir in dirnames[:]:
@@ -999,7 +1001,7 @@ def template(src, dest, user=None, group=None, mode=None, create_remote_dir=True
         line_number = trace_frames[-1][1]
 
         # Quickly read the line in question and one above/below for nicer debugging
-        with open(src, "r") as f:
+        with get_file_io(src, "r") as f:
             template_lines = f.readlines()
 
         template_lines = [line.strip() for line in template_lines]
@@ -1012,7 +1014,7 @@ def template(src, dest, user=None, group=None, mode=None, create_remote_dir=True
                 e,
                 "\n".join(relevant_lines),
             ),
-        )
+        ) from None
 
     output_file = StringIO(output)
     # Set the template attribute for nicer debugging