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

pyinfra/api/command.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import shlex
+from inspect import getfullargspec
+from string import Formatter
+from typing import IO, TYPE_CHECKING, Callable, Union
+
+import gevent
+from typing_extensions import Unpack, override
+
+from pyinfra.context import LocalContextObject, ctx_config, ctx_host
+
+from .arguments import ConnectorArguments
+
+if TYPE_CHECKING:
+    from pyinfra.api.host import Host
+    from pyinfra.api.state import State
+
+
+def make_formatted_string_command(string: str, *args, **kwargs) -> "StringCommand":
+    """
+    Helper function that takes a shell command or script as a string, splits it
+    using ``shlex.split`` and then formats each bit, returning a ``StringCommand``
+    instance with each bit.
+
+    Useful to enable string formatted commands/scripts, for example:
+
+    .. code:: python
+
+        curl_command = make_formatted_string_command(
+            'curl -sSLf {0} -o {1}',
+            QuoteString(src),
+            QuoteString(dest),
+        )
+    """
+
+    formatter = Formatter()
+    string_bits = []
+
+    for bit in shlex.split(string):
+        for item in formatter.parse(bit):
+            if item[0]:
+                string_bits.append(item[0])
+            if item[1]:
+                value, _ = formatter.get_field(item[1], args, kwargs)
+                string_bits.append(value)
+
+    return StringCommand(*string_bits)
+
+
+class MaskString(str):
+    pass
+
+
+class QuoteString:
+    obj: Union[str, "StringCommand"]
+
+    def __init__(self, obj: Union[str, "StringCommand"]):
+        self.obj = obj
+
+    @override
+    def __repr__(self) -> str:
+        return f"QuoteString({self.obj})"
+
+
+class PyinfraCommand:
+    connector_arguments: ConnectorArguments
+
+    def __init__(self, **arguments: Unpack[ConnectorArguments]):
+        self.connector_arguments = arguments
+
+    @override
+    def __eq__(self, other) -> bool:
+        if isinstance(other, self.__class__) and repr(self) == repr(other):
+            return True
+        return False
+
+    def execute(self, state: "State", host: "Host", connector_arguments: ConnectorArguments):
+        raise NotImplementedError
+
+
+class StringCommand(PyinfraCommand):
+    def __init__(
+        self,
+        *bits,
+        _separator=" ",
+        **arguments: Unpack[ConnectorArguments],
+    ):
+        super().__init__(**arguments)
+        self.bits = bits
+        self.separator = _separator
+
+    @override
+    def __str__(self) -> str:
+        return self.get_masked_value()
+
+    @override
+    def __repr__(self) -> str:
+        return f"StringCommand({self.get_masked_value()})"
+
+    def _get_all_bits(self, bit_accessor):
+        all_bits = []
+
+        for bit in self.bits:
+            quote = False
+            if isinstance(bit, QuoteString):
+                quote = True
+                bit = bit.obj
+
+            if isinstance(bit, StringCommand):
+                bit = bit_accessor(bit)
+
+            if not isinstance(bit, str):
+                bit = "{0}".format(bit)
+
+            if quote:
+                bit = shlex.quote(bit)
+
+            all_bits.append(bit)
+
+        return all_bits
+
+    def get_raw_value(self) -> str:
+        return self.separator.join(
+            self._get_all_bits(
+                lambda bit: bit.get_raw_value(),
+            ),
+        )
+
+    def get_masked_value(self) -> str:
+        return self.separator.join(
+            [
+                "***" if isinstance(bit, MaskString) else bit
+                for bit in self._get_all_bits(lambda bit: bit.get_masked_value())
+            ],
+        )
+
+    @override
+    def execute(self, state: "State", host: "Host", connector_arguments: ConnectorArguments):
+        connector_arguments.update(self.connector_arguments)
+
+        return host.run_shell_command(
+            self,
+            print_output=state.print_output,
+            print_input=state.print_input,
+            **connector_arguments,
+        )
+
+
+class FileUploadCommand(PyinfraCommand):
+    def __init__(
+        self,
+        src: str | IO,
+        dest: str,
+        remote_temp_filename=None,
+        **kwargs: Unpack[ConnectorArguments],
+    ):
+        super().__init__(**kwargs)
+        self.src = src
+        self.dest = dest
+        self.remote_temp_filename = remote_temp_filename
+
+    @override
+    def __repr__(self):
+        return "FileUploadCommand({0}, {1})".format(self.src, self.dest)
+
+    @override
+    def execute(self, state: "State", host: "Host", connector_arguments: ConnectorArguments):
+        connector_arguments.update(self.connector_arguments)
+
+        return host.put_file(
+            self.src,
+            self.dest,
+            remote_temp_filename=self.remote_temp_filename,
+            print_output=state.print_output,
+            print_input=state.print_input,
+            **connector_arguments,
+        )
+
+
+class FileDownloadCommand(PyinfraCommand):
+    def __init__(
+        self,
+        src: str,
+        dest: str | IO,
+        remote_temp_filename=None,
+        **kwargs: Unpack[ConnectorArguments],
+    ):
+        super().__init__(**kwargs)
+        self.src = src
+        self.dest = dest
+        self.remote_temp_filename = remote_temp_filename
+
+    @override
+    def __repr__(self):
+        return "FileDownloadCommand({0}, {1})".format(self.src, self.dest)
+
+    @override
+    def execute(self, state: "State", host: "Host", connector_arguments: ConnectorArguments):
+        connector_arguments.update(self.connector_arguments)
+
+        return host.get_file(
+            self.src,
+            self.dest,
+            remote_temp_filename=self.remote_temp_filename,
+            print_output=state.print_output,
+            print_input=state.print_input,
+            **connector_arguments,
+        )
+
+
+class FunctionCommand(PyinfraCommand):
+    def __init__(
+        self,
+        function: Callable,
+        args,
+        func_kwargs,
+        **kwargs: Unpack[ConnectorArguments],
+    ):
+        super().__init__(**kwargs)
+        self.function = function
+        self.args = args
+        self.kwargs = func_kwargs
+
+    @override
+    def __repr__(self):
+        return "FunctionCommand({0}, {1}, {2})".format(
+            self.function.__name__,
+            self.args,
+            self.kwargs,
+        )
+
+    @override
+    def execute(self, state: "State", host: "Host", connector_arguments: ConnectorArguments):
+        argspec = getfullargspec(self.function)
+        if "state" in argspec.args and "host" in argspec.args:
+            return self.function(state, host, *self.args, **self.kwargs)
+
+        # If we're already running inside a greenlet (ie a nested callback) just execute the func
+        # without any gevent.spawn which will break the local host object.
+        if isinstance(host, LocalContextObject):
+            self.function(*self.args, **self.kwargs)
+            return
+
+        def execute_function() -> None:
+            with ctx_config.use(state.config.copy()):
+                with ctx_host.use(host):
+                    self.function(*self.args, **self.kwargs)
+
+        greenlet = gevent.spawn(execute_function)
+        return greenlet.get()
+
+
+class RsyncCommand(PyinfraCommand):
+    def __init__(self, src: str, dest: str, flags, **kwargs: Unpack[ConnectorArguments]):
+        super().__init__(**kwargs)
+        self.src = src
+        self.dest = dest
+        self.flags = flags
+
+    @override
+    def __repr__(self):
+        return "RsyncCommand({0}, {1}, {2})".format(self.src, self.dest, self.flags)
+
+    @override
+    def execute(self, state: "State", host: "Host", connector_arguments: ConnectorArguments):
+        return host.rsync(
+            self.src,
+            self.dest,
+            self.flags,
+            print_output=state.print_output,
+            print_input=state.print_input,
+            **connector_arguments,
+        )

pyinfra/api/config.py

@@ -1,46 +1,240 @@
-class Config(object):
-    '''
-    The default/base configuration options for a pyinfra deploy.
-    '''
+try:
+    import importlib_metadata
+except ImportError:
+    import importlib.metadata as importlib_metadata  # type: ignore[no-redef]
 
-    state = None
+from os import path
+from typing import Iterable, Optional, Set
 
-    # % of hosts which have to fail for all operations to stop
-    FAIL_PERCENT = None
+from packaging.markers import Marker
+from packaging.requirements import Requirement
+from packaging.specifiers import SpecifierSet
+from packaging.version import Version
+from typing_extensions import override
 
-    # Seconds to timeout SSH connections
-    CONNECT_TIMEOUT = 10
+from pyinfra import __version__, state
+
+from .exceptions import PyinfraError
 
-    # Temporary directory (on the remote side) to use for caching any files/downloads
-    TEMP_DIR = '/tmp'
 
+class ConfigDefaults:
+    # % of hosts which have to fail for all operations to stop
+    FAIL_PERCENT: Optional[int] = None
+    # Seconds to timeout SSH connections
+    CONNECT_TIMEOUT: int = 10
+    # Temporary directory (on the remote side) to use for caching any files/downloads, the default
+    # None value first tries to load the hosts' temporary directory configured via "TMPDIR" env
+    # variable, falling back to DEFAULT_TEMP_DIR if not set.
+    TEMP_DIR: Optional[str] = None
+    DEFAULT_TEMP_DIR: str = "/tmp"
     # Gevent pool size (defaults to #of target hosts)
-    PARALLEL = None
+    PARALLEL: int = 0
+    # Specify the required pyinfra version (using PEP 440 setuptools specifier)
+    REQUIRE_PYINFRA_VERSION: Optional[str] = None
+    # Specify any required packages (either using PEP 440 or a requirements file)
+    # Note: this can also include pyinfra potentially replacing REQUIRE_PYINFRA_VERSION
+    REQUIRE_PACKAGES: Optional[str] = None
+    # All these can be overridden inside individual operation calls:
+    # Switch to this user (from ssh_user) using su before executing operations
+    SU_USER: Optional[str] = None
+    USE_SU_LOGIN: bool = False
+    SU_SHELL: bool = False
+    PRESERVE_SU_ENV: bool = False
+    # Use sudo and optional user
+    SUDO: bool = False
+    SUDO_USER: Optional[str] = None
+    PRESERVE_SUDO_ENV: bool = False
+    USE_SUDO_LOGIN: bool = False
+    SUDO_PASSWORD: Optional[str] = None
+    # Use doas and optional user
+    DOAS: bool = False
+    DOAS_USER: Optional[str] = None
+    # Only show errors but don't count as failure
+    IGNORE_ERRORS: bool = False
+    # Shell to use to execute commands
+    SHELL: str = "sh"
+    # Whether to display full diffs for files
+    DIFF: bool = False
+    # Number of times to retry failed operations
+    RETRY: int = 0
+    # Delay in seconds between retry attempts
+    RETRY_DELAY: int = 5
 
-    # Specify a minimum required pyinfra version for a deploy
-    MIN_PYINFRA_VERSION = None
 
-    # All these can be overridden inside individual operation calls:
+config_defaults = {key: value for key, value in ConfigDefaults.__dict__.items() if key.isupper()}
 
-    # Switch to this user (from ssh_user) using su before executing operations
-    SU_USER = None
 
-    # Use sudo and optional user
-    SUDO = False
-    SUDO_USER = None
-    PRESERVE_SUDO_ENV = False
+def check_pyinfra_version(version: str):
+    if not version:
+        return
+    running_version = Version(__version__)
+    required_versions = SpecifierSet(version)
 
-    # Only show errors, but don't count as failure
-    IGNORE_ERRORS = False
+    if running_version not in required_versions:
+        raise PyinfraError(
+            f"pyinfra version requirement not met (requires {version}, running {__version__})"
+        )
 
-    # Shell to use to execute commands
-    SHELL = 'sh'
+
+def _check_requirements(requirements: Iterable[str]) -> Set[Requirement]:
+    """
+    Check whether each of the given requirements and all their dependencies are
+    installed.
+
+    Or more precisely, this checks that each of the given *requirements* is
+    satisfied by some installed *distribution package*, and so on recursively
+    for each of the dependencies of those distribution packages. The terminology
+    here is as follows:
+
+    * A *distribution package* is essentially a thing that can be installed with
+      ``pip``, from an sdist or wheel or Git repo or so on.
+    * A *requirement* is the expectation that a distribution package satisfying
+      some constraint is installed.
+    * A *dependency* is a requirement specified by a distribution package (as
+      opposed to the requirements passed in to this function).
+
+    So what this function does is start from the given requirements, for each
+    one check that it is satisfied by some installed distribution package, and
+    if so recursively perform the same check on all the dependencies of that
+    distribution package. In short, it's traversing the graph of package
+    requirements. It stops whenever it finds a requirement that is not satisfied
+    (i.e. a required package that is not installed), or when it runs out of
+    requirements to check.
+
+    .. note::
+        This is basically equivalent to ``pkg_resources.require()`` except that
+        when ``require()`` succeeds, it will return the list of distribution
+        packages that satisfy the given requirements and their dependencies, and
+        when it fails, it will raise an exception. This function just returns
+        the requirements which were not satisfied instead.
+
+    :param requirements: The requirements to check for in the set of installed
+        packages (along with their dependencies).
+    :return: The set of requirements that were not satisfied, which will be
+        an empty set if all requirements (recursively) were satisfied.
+    """
+
+    # Based on pkg_resources.require() from setuptools. The implementation of
+    # hbutils.system.check_reqs() from the hbutils package was also helpful in
+    # clarifying what this is supposed to do.
+
+    reqs_to_check: Set[Requirement] = set(Requirement(r) for r in requirements)
+    reqs_satisfied: Set[Requirement] = set()
+    reqs_not_satisfied: Set[Requirement] = set()
+
+    while reqs_to_check:
+        req = reqs_to_check.pop()
+        assert req not in reqs_satisfied and req not in reqs_not_satisfied
+
+        # Check for an installed distribution package with the right name and version
+        try:
+            dist = importlib_metadata.distribution(req.name)
+        except importlib_metadata.PackageNotFoundError:
+            # No installed package with the right name
+            # This would raise a DistributionNotFound error from pkg_resources.require()
+            reqs_not_satisfied.add(req)
+            continue
+
+        if dist.version not in req.specifier:
+            # There is a distribution with the right name but wrong version
+            # This would raise a VersionConflict error from pkg_resources.require()
+            reqs_not_satisfied.add(req)
+            continue
+
+        reqs_satisfied.add(req)
+
+        # If the distribution package has dependencies of its own, go through
+        # those dependencies and for each one add it to the set to be checked if
+        # - it's unconditional (no marker)
+        # - or it's conditional and the condition is satisfied (the marker
+        #   evaluates to true) in the current environment
+        # Markers can check things like the Python version and system version
+        # etc., and/or they can check which extras of the distribution package
+        # were required. To facilitate checking extras we have to pass the extra
+        # in the environment when calling Marker.evaluate().
+        if dist.requires:
+            if req.extras:
+                extras_envs = [{"extra": extra} for extra in req.extras]
+
+                def evaluate_marker(marker: Marker) -> bool:
+                    return any(map(marker.evaluate, extras_envs))
+
+            else:
+
+                def evaluate_marker(marker: Marker) -> bool:
+                    return marker.evaluate()
+
+            for dist_req_str in dist.requires:
+                dist_req = Requirement(dist_req_str)
+                if dist_req in reqs_satisfied or dist_req in reqs_not_satisfied:
+                    continue
+                if (not dist_req.marker) or evaluate_marker(dist_req.marker):
+                    reqs_to_check.add(dist_req)
+
+    return reqs_not_satisfied
+
+
+def check_require_packages(requirements_config):
+    if not requirements_config:
+        return
+
+    if isinstance(requirements_config, (list, tuple)):
+        requirements = requirements_config
+    else:
+        with open(path.join(state.cwd or "", requirements_config), encoding="utf-8") as f:
+            requirements = [line.split("#egg=")[-1] for line in f.read().splitlines()]
+
+    requirements_not_met = _check_requirements(requirements)
+    if requirements_not_met:
+        raise PyinfraError(
+            "Deploy requirements ({0}) not met: missing {1}".format(
+                requirements_config, ", ".join(str(r) for r in requirements_not_met)
+            )
+        )
+
+
+config_checkers = {
+    "REQUIRE_PYINFRA_VERSION": check_pyinfra_version,
+    "REQUIRE_PACKAGES": check_require_packages,
+}
+
+
+class Config(ConfigDefaults):
+    """
+    The default/base configuration options for a pyinfra deploy.
+    """
 
     def __init__(self, **kwargs):
         # Always apply some env
-        env = kwargs.pop('ENV', {})
+        env = kwargs.pop("ENV", {})
         self.ENV = env
 
-        # Apply kwargs
-        for key, value in kwargs.items():
+        config = config_defaults.copy()
+        config.update(kwargs)
+
+        for key, value in config.items():
+            setattr(self, key, value)
+
+    @override
+    def __setattr__(self, key, value):
+        super().__setattr__(key, value)
+
+        checker = config_checkers.get(key)
+        if checker:
+            checker(value)
+
+    def get_current_state(self):
+        return [(key, getattr(self, key)) for key in config_defaults.keys()]
+
+    def set_current_state(self, config_state):
+        for key, value in config_state:
             setattr(self, key, value)
+
+    def lock_current_state(self) -> None:
+        self._locked_config = self.get_current_state()
+
+    def reset_locked_state(self) -> None:
+        self.set_current_state(self._locked_config)
+
+    def copy(self) -> "Config":
+        return Config(**dict(self.get_current_state()))

pyinfra/api/connect.py

@@ -1,25 +1,28 @@
+from typing import TYPE_CHECKING
+
 import gevent
 
 from pyinfra.progress import progress_spinner
 
+if TYPE_CHECKING:
+    from pyinfra.api.state import State
+
 
-def connect_all(state):
-    '''
+def connect_all(state: "State"):
+    """
     Connect to all the configured servers in parallel. Reads/writes state.inventory.
 
     Args:
         state (``pyinfra.api.State`` obj): the state containing an inventory to connect to
-    '''
+    """
 
     hosts = [
-        host for host in state.inventory
-        if state.is_host_in_limit(host)
+        host
+        for host in state.inventory
+        if state.is_host_in_limit(host)  # these are the hosts to activate ("initially connect to")
     ]
 
-    greenlet_to_host = {
-        state.pool.spawn(host.connect, state): host
-        for host in hosts
-    }
+    greenlet_to_host = {state.pool.spawn(host.connect): host for host in hosts}
 
     with progress_spinner(greenlet_to_host.values()) as progress:
         for greenlet in gevent.iwait(greenlet_to_host.keys()):
@@ -33,7 +36,7 @@ def connect_all(state):
         # Raise any unexpected exception
         greenlet.get()
 
-        if host.connection:
+        if host.connected:
             state.activate_host(host)
         else:
             failed_hosts.add(host)
@@ -42,6 +45,23 @@ def connect_all(state):
     state.fail_hosts(failed_hosts, activated_count=len(hosts))
 
 
-def disconnect_all(state):
-    for host in state.activated_hosts:  # only hosts we connected to please!
-        host.disconnect(state)  # normally a noop
+def disconnect_all(state: "State"):
+    """
+    Disconnect from all of the configured servers in parallel. Reads/writes state.inventory.
+
+    Args:
+        state (``pyinfra.api.State`` obj): the state containing an inventory to connect to
+    """
+    greenlet_to_host = {
+        state.pool.spawn(host.disconnect): host
+        for host in state.activated_hosts  # only hosts we connected to please!
+    }
+
+    with progress_spinner(greenlet_to_host.values()) as progress:
+        for greenlet in gevent.iwait(greenlet_to_host.keys()):
+            host = greenlet_to_host[greenlet]
+            progress(host)
+
+    for greenlet, host in greenlet_to_host.items():
+        # Raise any unexpected exception
+        greenlet.get()

pyinfra/api/connectors.py

@@ -0,0 +1,27 @@
+try:
+    from importlib_metadata import entry_points
+except ImportError:
+    from importlib.metadata import entry_points  # type: ignore[assignment]
+
+
+def _load_connector(entrypoint):
+    return entrypoint.load()
+
+
+def get_all_connectors():
+    return {
+        entrypoint.name: _load_connector(entrypoint)
+        for entrypoint in entry_points(group="pyinfra.connectors")
+    }
+
+
+def get_execution_connectors():
+    return {
+        connector: connector_mod
+        for connector, connector_mod in get_all_connectors().items()
+        if connector_mod.handles_execution
+    }
+
+
+def get_execution_connector(name):
+    return get_execution_connectors()[name]