pyinfra 3.0.dev0__py2.py3-none-any.whl → 3.0.1__py2.py3-none-any.whl

pyinfra/connectors/base.py

@@ -1,42 +1,90 @@
 from __future__ import annotations
 
 import abc
-from dataclasses import dataclass
 from io import IOBase
-from typing import TYPE_CHECKING, Iterable, Iterator, Optional, Union
-
-from typing_extensions import Unpack
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    Iterator,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    get_type_hints,
+)
+
+from typing_extensions import TypedDict, Unpack
+
+from pyinfra.api.exceptions import ConnectorDataTypeError
+from pyinfra.api.util import raise_if_bad_type
 
 if TYPE_CHECKING:
     from pyinfra.api.arguments import ConnectorArguments
     from pyinfra.api.command import StringCommand
-    from pyinfra.api.host import Host
+    from pyinfra.api.host import Host, HostData
     from pyinfra.api.state import State
 
     from .util import CommandOutput
 
 
-def make_keys(prefix: str, cls):
-    class Keys:
-        pass
+T = TypeVar("T")
+default_sentinel = object()
+
+
+def host_to_connector_data(
+    connector_data: Type[T],
+    connector_data_meta: dict[str, DataMeta],
+    host_data: "HostData",
+) -> T:
+    data: T = cast(T, {})
+    for key, type_ in get_type_hints(connector_data).items():
+        value = host_data.get(key, default_sentinel)
+        if value is default_sentinel:
+            value = connector_data_meta[key].default
+        else:
+            raise_if_bad_type(
+                value,
+                type_,
+                ConnectorDataTypeError,
+                f"Invalid connector data `{key}`:",
+            )
+
+        data[key] = value  # type: ignore
+    return data
+
+
+class DataMeta:
+    description: str
+    default: Any
+
+    def __init__(self, description, default=None) -> None:
+        self.description = description
+        self.default = default
 
-    for key in cls.__dict__:
-        if not key.startswith("_"):
-            setattr(Keys, key, f"{prefix}_{key}")
 
-    return Keys
+class ConnectorData(TypedDict, total=False):
+    pass
 
 
-@dataclass
 class BaseConnector(abc.ABC):
     state: "State"
     host: "Host"
 
     handles_execution = False
 
+    data_cls: Type = ConnectorData
+    data_meta: dict[str, DataMeta] = {}
+
+    def __init__(self, state: "State", host: "Host"):
+        self.state = state
+        self.host = host
+        self.data = host_to_connector_data(self.data_cls, self.data_meta, host.data)
+
     @staticmethod
     @abc.abstractmethod
-    def make_names_data(id: str) -> Iterator[tuple[str, dict, list[str]]]:
+    def make_names_data(name: str) -> Iterator[tuple[str, dict, list[str]]]:
         """
         Generates hosts/data/groups information for inventory. This allows a
         single connector reference to generate multiple target hosts.
@@ -60,8 +108,7 @@ class BaseConnector(abc.ABC):
         print_output: bool,
         print_input: bool,
         **arguments: Unpack["ConnectorArguments"],
-    ) -> tuple[bool, "CommandOutput"]:
-        ...
+    ) -> tuple[bool, "CommandOutput"]: ...
 
     @abc.abstractmethod
     def put_file(
@@ -72,8 +119,7 @@ class BaseConnector(abc.ABC):
         print_output: bool = False,
         print_input: bool = False,
         **arguments: Unpack["ConnectorArguments"],
-    ) -> bool:
-        ...
+    ) -> bool: ...
 
     @abc.abstractmethod
     def get_file(
@@ -84,8 +130,7 @@ class BaseConnector(abc.ABC):
         print_output: bool = False,
         print_input: bool = False,
         **arguments: Unpack["ConnectorArguments"],
-    ) -> bool:
-        ...
+    ) -> bool: ...
 
     def check_can_rsync(self):
         raise NotImplementedError("This connector does not support rsync")

pyinfra/connectors/chroot.py

@@ -27,6 +27,10 @@ def show_warning():
 
 
 class ChrootConnector(BaseConnector):
+    """
+    The chroot connector allows you to execute operations within another root.
+    """
+
     handles_execution = True
 
     local: LocalConnector
@@ -36,17 +40,17 @@ class ChrootConnector(BaseConnector):
         self.local = LocalConnector(state, host)
 
     @staticmethod
-    def make_names_data(directory: Optional[str] = None):
-        if not directory:
+    def make_names_data(name: Optional[str] = None):
+        if not name:
             raise InventoryError("No directory provided!")
 
         show_warning()
 
-        yield "@chroot/{0}".format(directory), {
-            "chroot_directory": "/{0}".format(directory.lstrip("/")),
+        yield "@chroot/{0}".format(name), {
+            "chroot_directory": "/{0}".format(name.lstrip("/")),
         }, ["@chroot"]
 
-    def connect(self):
+    def connect(self) -> None:
         self.local.connect()
 
         chroot_directory = self.host.data.chroot_directory
@@ -61,7 +65,6 @@ class ChrootConnector(BaseConnector):
             raise ConnectError(e.args[0])
 
         self.host.connector_data["chroot_directory"] = chroot_directory
-        return True
 
     def run_shell_command(
         self,
@@ -117,11 +120,10 @@ class ChrootConnector(BaseConnector):
                     temp_f.write(data)
 
             chroot_directory = self.host.connector_data["chroot_directory"]
-
-            chroot_command = "cp {0} {1}/{2}".format(
+            chroot_command = StringCommand(
+                "cp",
                 temp_filename,
-                chroot_directory,
-                remote_filename,
+                f"{chroot_directory}/{remote_filename}",
             )
 
             status, output = self.local.run_shell_command(
@@ -159,9 +161,9 @@ class ChrootConnector(BaseConnector):
 
         try:
             chroot_directory = self.host.connector_data["chroot_directory"]
-            chroot_command = "cp {0}/{1} {2}".format(
-                chroot_directory,
-                remote_filename,
+            chroot_command = StringCommand(
+                "cp",
+                f"{chroot_directory}/{remote_filename}",
                 temp_filename,
             )
 

pyinfra/connectors/docker.py

@@ -1,26 +1,3 @@
-"""
-The ``@docker`` connector allows you to build Docker images, or modify running
-Docker containers, using ``pyinfra``. You can pass either an image name or
-existing container ID:
-
-+ Image - will create a container from the image, execute operations and save
-    into a new image
-+ Existing container ID - will simply execute operations against the container,
-    leaving it up afterwards
-
-
-.. code:: shell
-
-    # A Docker base image must be provided
-    pyinfra @docker/alpine:3.8 ...
-
-    # pyinfra can run on multiple Docker images in parallel
-    pyinfra @docker/alpine:3.8,@docker/ubuntu:bionic ...
-
-    # Execute against a running container
-    pyinfra @docker/2beb8c15a1b1 ...
-"""
-
 from __future__ import annotations
 
 import json
@@ -29,7 +6,7 @@ from tempfile import mkstemp
 from typing import TYPE_CHECKING
 
 import click
-from typing_extensions import Unpack
+from typing_extensions import TypedDict, Unpack
 
 from pyinfra import local, logger
 from pyinfra.api import QuoteString, StringCommand
@@ -37,7 +14,7 @@ from pyinfra.api.exceptions import ConnectError, InventoryError, PyinfraError
 from pyinfra.api.util import get_file_io
 from pyinfra.progress import progress_spinner
 
-from .base import BaseConnector, make_keys
+from .base import BaseConnector, DataMeta
 from .local import LocalConnector
 from .util import CommandOutput, extract_control_arguments, make_unix_command_for_host
 
@@ -47,20 +24,24 @@ if TYPE_CHECKING:
     from pyinfra.api.state import State
 
 
-class DataKeys:
-    identifier = "ID of container or image to target"
-    container_id = "ID of container to target, overrides ``docker_identifier``"
+class ConnectorData(TypedDict):
+    docker_identifier: str
 
 
-DATA_KEYS = make_keys("docker", DataKeys)
+connector_data_meta: dict[str, DataMeta] = {
+    "docker_identifier": DataMeta("ID of container or image to start from"),
+}
 
 
-def _find_start_docker_container(container_id):
+def _find_start_docker_container(container_id) -> tuple[str, bool]:
     docker_info = local.shell("docker container inspect {0}".format(container_id))
+    assert isinstance(docker_info, str)
     docker_info = json.loads(docker_info)[0]
     if docker_info["State"]["Running"] is False:
         logger.info("Starting stopped container: {0}".format(container_id))
         local.shell("docker container start {0}".format(container_id))
+        return container_id, False
+    return container_id, True
 
 
 def _start_docker_image(image_name):
@@ -76,53 +57,69 @@ def _start_docker_image(image_name):
 
 
 class DockerConnector(BaseConnector):
+    """
+    The docker connector allows you to build Docker images or modify running
+    Docker containers. You can pass either an image name or existing container ID:
+
+    + Image - will create a new container from the image, execute operations \
+        against it, save into a new Docker image and remove the container
+    + Existing container ID - will execute operations against the running \
+        container, leaving it running
+
+    .. code:: shell
+
+        # A Docker base image must be provided
+        pyinfra @docker/alpine:3.8 ...
+
+        # pyinfra can run on multiple Docker images in parallel
+        pyinfra @docker/alpine:3.8,@docker/ubuntu:bionic ...
+
+        # Execute against a running container
+        pyinfra @docker/2beb8c15a1b1 ...
+    """
+
     handles_execution = True
-    data_keys = DATA_KEYS
+
+    data_cls = ConnectorData
+    data_meta = connector_data_meta
+    data: ConnectorData
 
     local: LocalConnector
 
+    container_id: str
+    no_stop: bool = False
+
     def __init__(self, state: "State", host: "Host"):
         super().__init__(state, host)
         self.local = LocalConnector(state, host)
 
     @staticmethod
-    def make_names_data(identifier=None):
-        if not identifier:
+    def make_names_data(name=None):
+        if not name:
             raise InventoryError("No docker base ID provided!")
 
         yield (
-            "@docker/{0}".format(identifier),
-            {DATA_KEYS.identifier: identifier},
+            "@docker/{0}".format(name),
+            {"docker_identifier": name},
             ["@docker"],
         )
 
-    def connect(self):
+    def connect(self) -> None:
         self.local.connect()
 
-        docker_container_id = self.host.data.get(DATA_KEYS.container_id)
-        if docker_container_id:  # user can provide a docker_container_id
-            self.host.connector_data["docker_container_no_disconnect"] = True
-            self.host.connector_data["docker_container_id"] = docker_container_id
-            return True
-
-        docker_identifier = getattr(self.host.data, DATA_KEYS.identifier)
+        docker_identifier = self.data["docker_identifier"]
         with progress_spinner({"prepare docker container"}):
             try:
-                # Check if the provided @docker/X is an existing container ID
-                _find_start_docker_container(docker_identifier)
+                self.container_id, was_running = _find_start_docker_container(docker_identifier)
+                if was_running:
+                    self.no_stop = True
             except PyinfraError:
-                container_id = _start_docker_image(docker_identifier)
-            else:
-                container_id = docker_identifier
-                self.host.connector_data["docker_container_no_disconnect"] = True
-
-        self.host.connector_data["docker_container_id"] = container_id
-        return True
+                self.container_id = _start_docker_image(docker_identifier)
 
     def disconnect(self):
-        container_id = self.host.connector_data["docker_container_id"]
+        container_id = self.container_id
 
-        if self.host.connector_data.get("docker_container_no_disconnect"):
+        if self.no_stop:
             logger.info(
                 "{0}docker build complete, container left running: {1}".format(
                     self.host.print_prefix,
@@ -157,7 +154,7 @@ class DockerConnector(BaseConnector):
     ) -> tuple[bool, CommandOutput]:
         local_arguments = extract_control_arguments(arguments)
 
-        container_id = self.host.connector_data["docker_container_id"]
+        container_id = self.container_id
 
         command = make_unix_command_for_host(self.state, self.host, command, **arguments)
         command = StringCommand(QuoteString(command))
@@ -207,11 +204,11 @@ class DockerConnector(BaseConnector):
 
                     temp_f.write(data)
 
-            docker_id = self.host.connector_data["docker_container_id"]
-            docker_command = "docker cp {0} {1}:{2}".format(
+            docker_command = StringCommand(
+                "docker",
+                "cp",
                 temp_filename,
-                docker_id,
-                remote_filename,
+                f"{self.container_id}:{remote_filename}",
             )
 
             status, output = self.local.run_shell_command(
@@ -254,10 +251,10 @@ class DockerConnector(BaseConnector):
         fd, temp_filename = mkstemp()
 
         try:
-            docker_id = self.host.connector_data["docker_container_id"]
-            docker_command = "docker cp {0}:{1} {2}".format(
-                docker_id,
-                remote_filename,
+            docker_command = StringCommand(
+                "docker",
+                "cp",
+                f"{self.container_id}:{remote_filename}",
                 temp_filename,
             )
 
@@ -268,17 +265,10 @@ class DockerConnector(BaseConnector):
             )
 
             # Load the temporary file and write it to our file or IO object
-            with open(temp_filename, encoding="utf-8") as temp_f:
+            with open(temp_filename, "rb") as temp_f:
                 with get_file_io(filename_or_io, "wb") as file_io:
                     data = temp_f.read()
-                    data_bytes: bytes
-
-                    if isinstance(data, str):
-                        data_bytes = data.encode()
-                    else:
-                        data_bytes = data
-
-                    file_io.write(data_bytes)
+                    file_io.write(data)
         finally:
             os.close(fd)
             os.remove(temp_filename)

pyinfra/connectors/dockerssh.py

@@ -1,17 +1,3 @@
-"""
-**Note**: this connector is in beta!
-
-The ``@dockerssh`` connector allows you to run commands on Docker containers on a remote machine.
-
-.. code:: shell
-
-    # A Docker base image must be provided
-    pyinfra @dockerssh/remotehost:alpine:3.8 ...
-
-    # pyinfra can run on multiple Docker images in parallel
-    pyinfra @dockerssh/remotehost:alpine:3.8,@dockerssh/remotehost:ubuntu:bionic ...
-"""
-
 import os
 from tempfile import mkstemp
 from typing import TYPE_CHECKING
@@ -41,6 +27,21 @@ def show_warning():
 
 
 class DockerSSHConnector(BaseConnector):
+    """
+    **Note**: this connector is in beta!
+
+    The ``@dockerssh`` connector allows you to run commands on Docker containers \
+    on a remote machine.
+
+    .. code:: shell
+
+        # A Docker base image must be provided
+        pyinfra @dockerssh/remotehost:alpine:3.8 ...
+
+        # pyinfra can run on multiple Docker images in parallel
+        pyinfra @dockerssh/remotehost:alpine:3.8,@dockerssh/remotehost:ubuntu:bionic ...
+    """
+
     handles_execution = True
 
     ssh: SSHConnector
@@ -50,10 +51,10 @@ class DockerSSHConnector(BaseConnector):
         self.ssh = SSHConnector(state, host)
 
     @staticmethod
-    def make_names_data(host_image_str):
+    def make_names_data(name):
         try:
-            hostname, image = host_image_str.split(":", 1)
-        except (AttributeError, ValueError):  # failure to parse the host_image_str
+            hostname, image = name.split(":", 1)
+        except (AttributeError, ValueError):  # failure to parse the name
             raise InventoryError("No ssh host or docker base image provided!")
 
         if not image:
@@ -164,7 +165,7 @@ class DockerSSHConnector(BaseConnector):
         """
 
         fd, local_temp_filename = mkstemp()
-        remote_temp_filename = remote_temp_filename or self.state.get_temp_filename(
+        remote_temp_filename = remote_temp_filename or self.host.get_temp_filename(
             local_temp_filename
         )
 
@@ -234,7 +235,7 @@ class DockerSSHConnector(BaseConnector):
         location and then reading that into our final file/IO object.
         """
 
-        remote_temp_filename = remote_temp_filename or self.state.get_temp_filename(remote_filename)
+        remote_temp_filename = remote_temp_filename or self.host.get_temp_filename(remote_filename)
 
         try:
             docker_id = self.host.host_data["docker_container_id"]

pyinfra/connectors/local.py

@@ -1,9 +1,5 @@
-"""
-The ``@local`` connector executes changes on the local machine using subprocesses.
-"""
-
 import os
-from distutils.spawn import find_executable
+from shutil import which
 from tempfile import mkstemp
 from typing import TYPE_CHECKING, Tuple
 
@@ -28,18 +24,30 @@ 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**.
+
+    Examples:
+
+    .. code::
+
+        # Install nginx
+        pyinfra inventory.py apt.packages nginx update=true _sudo=true
+    """
+
     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"]
 
     def run_shell_command(
         self,
-        command,
+        command: StringCommand,
         print_output: bool = False,
         print_input: bool = False,
         **arguments: Unpack["ConnectorArguments"],
@@ -48,17 +56,14 @@ class LocalConnector(BaseConnector):
         Execute a command on the local machine.
 
         Args:
-            state (``pyinfra.api.State`` object): state object for this command
-            host (``pyinfra.api.Host`` object): the target host
-            command (string): actual command to execute
-            sudo (boolean): whether to wrap the command with sudo
-            sudo_user (string): user to sudo to
-            env (dict): environment variables to set
-            timeout (int): timeout for this command to complete before erroring
+            command (StringCommand): actual command to execute
+            print_output (bool): whether to print command output
+            print_input (bool): whether to print command input
+            arguments: (ConnectorArguments): connector global arguments
 
         Returns:
-            tuple: (exit_code, stdout, stderr)
-            stdout and stderr are both lists of strings from each buffer.
+            tuple: (bool, CommandOutput)
+            Bool indicating success and CommandOutput with stdout/stderr lines.
         """
 
         arguments.pop("_get_pty", False)
@@ -108,6 +113,9 @@ class LocalConnector(BaseConnector):
         """
         Upload a local file or IO object by copying it to a temporary directory
         and then writing it to the upload location.
+
+        Returns:
+            bool: Indicating success or failure
         """
 
         _, temp_filename = mkstemp()
@@ -156,6 +164,9 @@ class LocalConnector(BaseConnector):
         """
         Download a local file by copying it to a temporary location and then writing
         it to our filename or IO object.
+
+        Returns:
+            bool: Indicating success or failure
         """
 
         _, temp_filename = mkstemp()
@@ -163,7 +174,7 @@ class LocalConnector(BaseConnector):
         try:
             # Copy the file using `cp` such that we support sudo/su
             status, output = self.run_shell_command(
-                "cp {0} {1}".format(remote_filename, temp_filename),
+                StringCommand("cp", remote_filename, temp_filename),
                 print_output=print_output,
                 print_input=print_input,
                 **arguments,
@@ -195,9 +206,8 @@ class LocalConnector(BaseConnector):
 
         return True
 
-    @staticmethod
-    def check_can_rsync(host):
-        if not find_executable("rsync"):
+    def check_can_rsync(self):
+        if not which("rsync"):
             raise NotImplementedError("The `rsync` binary is not available on this system.")
 
     def rsync(
@@ -210,7 +220,7 @@ class LocalConnector(BaseConnector):
         **arguments,
     ) -> bool:
         status, output = self.run_shell_command(
-            "rsync {0} {1} {2}".format(" ".join(flags), src, dest),
+            StringCommand("rsync", " ".join(flags), src, dest),
             print_output=print_output,
             print_input=print_input,
             **arguments,