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

pyinfra/connectors/base.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import abc
+from io import IOBase
+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, HostData
+    from pyinfra.api.state import State
+
+    from .util import CommandOutput
+
+
+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
+
+
+class ConnectorData(TypedDict, total=False):
+    pass
+
+
+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(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.
+        """
+        ...
+
+    def connect(self) -> None:
+        """
+        Connect this connector instance.
+        """
+
+    def disconnect(self) -> None:
+        """
+        Disconnect this connector instance.
+        """
+
+    @abc.abstractmethod
+    def run_shell_command(
+        self,
+        command: "StringCommand",
+        print_output: bool,
+        print_input: bool,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> tuple[bool, "CommandOutput"]: ...
+
+    @abc.abstractmethod
+    def put_file(
+        self,
+        filename_or_io: Union[str, IOBase],
+        remote_filename: str,
+        remote_temp_filename: Optional[str] = None,
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> bool: ...
+
+    @abc.abstractmethod
+    def get_file(
+        self,
+        remote_filename: str,
+        filename_or_io: Union[str, IOBase],
+        remote_temp_filename: Optional[str] = None,
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> bool: ...
+
+    def check_can_rsync(self):
+        raise NotImplementedError("This connector does not support rsync")
+
+    def rsync(
+        self,
+        src: str,
+        dest: str,
+        flags: Iterable[str],
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> bool:
+        raise NotImplementedError("This connector does not support rsync")

pyinfra/connectors/chroot.py

@@ -3,210 +3,201 @@ from tempfile import mkstemp
 from typing import TYPE_CHECKING, Optional
 
 import click
+from typing_extensions import Unpack
 
 from pyinfra import local, logger
 from pyinfra.api import QuoteString, StringCommand
-from pyinfra.api.connectors import BaseConnectorMeta
 from pyinfra.api.exceptions import ConnectError, InventoryError, PyinfraError
 from pyinfra.api.util import get_file_io, memoize
 from pyinfra.progress import progress_spinner
 
-from .local import run_shell_command as run_local_shell_command
-from .util import make_unix_command_for_host
+from .base import BaseConnector
+from .local import LocalConnector
+from .util import extract_control_arguments, make_unix_command_for_host
 
 if TYPE_CHECKING:
+    from pyinfra.api.arguments import ConnectorArguments
     from pyinfra.api.host import Host
     from pyinfra.api.state import State
 
 
-class Meta(BaseConnectorMeta):
-    handles_execution = True
-
-
 @memoize
 def show_warning():
     logger.warning("The @chroot connector is in beta!")
 
 
-def make_names_data(directory: Optional[str] = None):
-    if not directory:
-        raise InventoryError("No directory provided!")
+class ChrootConnector(BaseConnector):
+    """
+    The chroot connector allows you to execute operations within another root.
+    """
 
-    show_warning()
+    handles_execution = True
 
-    yield "@chroot/{0}".format(directory), {
-        "chroot_directory": "/{0}".format(directory.lstrip("/")),
-    }, ["@chroot"]
+    local: LocalConnector
 
+    def __init__(self, state: "State", host: "Host"):
+        super().__init__(state, host)
+        self.local = LocalConnector(state, host)
 
-def connect(state: "State", host: "Host"):
-    chroot_directory = host.data.chroot_directory
+    @staticmethod
+    def make_names_data(name: Optional[str] = None):
+        if not name:
+            raise InventoryError("No directory provided!")
 
-    try:
-        with progress_spinner({"chroot run"}):
-            local.shell(
-                "chroot {0} ls".format(chroot_directory),
-                splitlines=True,
-            )
-    except PyinfraError as e:
-        raise ConnectError(e.args[0])
-
-    host.connector_data["chroot_directory"] = chroot_directory
-    return True
-
-
-def run_shell_command(
-    state: "State",
-    host: "Host",
-    command,
-    get_pty: bool = False,
-    timeout=None,
-    stdin=None,
-    success_exit_codes=None,
-    print_output: bool = False,
-    print_input: bool = False,
-    return_combined_output: bool = False,
-    **command_kwargs,
-):
-    chroot_directory = host.connector_data["chroot_directory"]
-
-    command = make_unix_command_for_host(state, host, command, **command_kwargs)
-    command = QuoteString(command)
-
-    logger.debug("--> Running chroot command on (%s): %s", chroot_directory, command)
-
-    chroot_command = StringCommand(
-        "chroot",
-        chroot_directory,
-        "sh",
-        "-c",
-        command,
-    )
-
-    return run_local_shell_command(
-        state,
-        host,
-        chroot_command,
-        timeout=timeout,
-        stdin=stdin,
-        success_exit_codes=success_exit_codes,
-        print_output=print_output,
-        print_input=print_input,
-        return_combined_output=return_combined_output,
-    )
-
-
-def put_file(
-    state: "State",
-    host: "Host",
-    filename_or_io,
-    remote_filename,
-    remote_temp_filename=None,  # ignored
-    print_output: bool = False,
-    print_input: bool = False,
-    **kwargs,  # ignored (sudo/etc)
-):
-
-    _, temp_filename = mkstemp()
-
-    try:
-        # Load our file or IO object and write it to the temporary file
-        with get_file_io(filename_or_io) as file_io:
-            with open(temp_filename, "wb") as temp_f:
-                data = file_io.read()
-
-                if isinstance(data, str):
-                    data = data.encode()
-
-                temp_f.write(data)
-
-        chroot_directory = host.connector_data["chroot_directory"]
-
-        chroot_command = "cp {0} {1}/{2}".format(
-            temp_filename,
-            chroot_directory,
-            remote_filename,
-        )
+        show_warning()
 
-        status, _, stderr = run_local_shell_command(
-            state,
-            host,
-            chroot_command,
-            print_output=print_output,
-            print_input=print_input,
-        )
-    finally:
-        os.remove(temp_filename)
-
-    if not status:
-        raise IOError("\n".join(stderr))
-
-    if print_output:
-        click.echo(
-            "{0}file uploaded to chroot: {1}".format(
-                host.print_prefix,
-                remote_filename,
-            ),
-            err=True,
-        )
+        yield "@chroot/{0}".format(name), {
+            "chroot_directory": "/{0}".format(name.lstrip("/")),
+        }, ["@chroot"]
+
+    def connect(self) -> None:
+        self.local.connect()
 
-    return status
+        chroot_directory = self.host.data.chroot_directory
 
+        try:
+            with progress_spinner({"chroot run"}):
+                local.shell(
+                    "chroot {0} ls".format(chroot_directory),
+                    splitlines=True,
+                )
+        except PyinfraError as e:
+            raise ConnectError(e.args[0])
 
-def get_file(
-    state: "State",
-    host: "Host",
-    remote_filename,
-    filename_or_io,
-    remote_temp_filename=None,  # ignored
-    print_output: bool = False,
-    print_input: bool = False,
-    **kwargs,  # ignored (sudo/etc)
-):
+        self.host.connector_data["chroot_directory"] = chroot_directory
 
-    _, temp_filename = mkstemp()
+    def run_shell_command(
+        self,
+        command,
+        print_output: bool = False,
+        print_input: bool = False,
+        **command_arguments: Unpack["ConnectorArguments"],
+    ):
+        local_arguments = extract_control_arguments(command_arguments)
 
-    try:
-        chroot_directory = host.connector_data["chroot_directory"]
-        chroot_command = "cp {0}/{1} {2}".format(
+        chroot_directory = self.host.connector_data["chroot_directory"]
+
+        command = make_unix_command_for_host(self.state, self.host, command, **command_arguments)
+        command = QuoteString(command)
+
+        logger.debug("--> Running chroot command on (%s): %s", chroot_directory, command)
+
+        chroot_command = StringCommand(
+            "chroot",
             chroot_directory,
-            remote_filename,
-            temp_filename,
+            "sh",
+            "-c",
+            command,
         )
 
-        status, _, stderr = run_local_shell_command(
-            state,
-            host,
+        return self.local.run_shell_command(
             chroot_command,
             print_output=print_output,
             print_input=print_input,
+            **local_arguments,
         )
 
-        # Load the temporary file and write it to our file or IO object
-        with open(temp_filename, encoding="utf-8") 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)
-    finally:
-        os.remove(temp_filename)
-
-    if not status:
-        raise IOError("\n".join(stderr))
-
-    if print_output:
-        click.echo(
-            "{0}file downloaded from chroot: {1}".format(
-                host.print_prefix,
-                remote_filename,
-            ),
-            err=True,
-        )
+    def put_file(
+        self,
+        filename_or_io,
+        remote_filename,
+        remote_temp_filename=None,  # ignored
+        print_output: bool = False,
+        print_input: bool = False,
+        **kwargs,  # ignored (sudo/etc)
+    ):
+        _, temp_filename = mkstemp()
+
+        try:
+            # Load our file or IO object and write it to the temporary file
+            with get_file_io(filename_or_io) as file_io:
+                with open(temp_filename, "wb") as temp_f:
+                    data = file_io.read()
+
+                    if isinstance(data, str):
+                        data = data.encode()
+
+                    temp_f.write(data)
+
+            chroot_directory = self.host.connector_data["chroot_directory"]
+            chroot_command = StringCommand(
+                "cp",
+                temp_filename,
+                f"{chroot_directory}/{remote_filename}",
+            )
+
+            status, output = self.local.run_shell_command(
+                chroot_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+        finally:
+            os.remove(temp_filename)
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file uploaded to chroot: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
+
+        return status
+
+    def get_file(
+        self,
+        remote_filename,
+        filename_or_io,
+        remote_temp_filename=None,  # ignored
+        print_output: bool = False,
+        print_input: bool = False,
+        **kwargs,  # ignored (sudo/etc)
+    ):
+        _, temp_filename = mkstemp()
+
+        try:
+            chroot_directory = self.host.connector_data["chroot_directory"]
+            chroot_command = StringCommand(
+                "cp",
+                f"{chroot_directory}/{remote_filename}",
+                temp_filename,
+            )
+
+            status, output = self.local.run_shell_command(
+                chroot_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+
+            # Load the temporary file and write it to our file or IO object
+            with open(temp_filename, encoding="utf-8") 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)
+        finally:
+            os.remove(temp_filename)
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file downloaded from chroot: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
 
-    return status
+        return status