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

pyinfra/operations/postgresql.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from pyinfra.api import operation
 
 from . import postgres

pyinfra/operations/puppet.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 from pyinfra.api import operation
 
 
 @operation(is_idempotent=False)
-def agent(server=None, port=None):
+def agent(server: str | None = None, port: int | None = None):
     """
     Run puppet agent
 

pyinfra/operations/python.py

@@ -2,11 +2,15 @@
 The Python module allows you to execute Python code within the context of a deploy.
 """
 
+from __future__ import annotations
+
+from typing import Callable
+
 from pyinfra.api import FunctionCommand, operation
 
 
 @operation(is_idempotent=False, _set_in_op=False)
-def call(function, *args, **kwargs):
+def call(function: Callable, *args, **kwargs):
     """
     Execute a Python function within a deploy.
 
@@ -43,7 +47,7 @@ def call(function, *args, **kwargs):
 
 
 @operation(is_idempotent=False, _set_in_op=False)
-def raise_exception(exception, *args, **kwargs):
+def raise_exception(exception: Exception, *args, **kwargs):
     """
     Raise a Python exception within a deploy.
 
@@ -63,6 +67,6 @@ def raise_exception(exception, *args, **kwargs):
     """
 
     def raise_exc(*args, **kwargs):  # pragma: no cover
-        raise exception(*args, **kwargs)
+        raise exception(*args, **kwargs)  # type: ignore[operator]
 
     yield FunctionCommand(raise_exc, args, kwargs)

pyinfra/operations/runit.py

@@ -0,0 +1,182 @@
+"""
+Manage runit services.
+"""
+
+from typing import Optional
+
+from pyinfra import host
+from pyinfra.api import operation
+from pyinfra.facts.files import File
+from pyinfra.facts.runit import RunitManaged, RunitStatus
+
+from .files import file, link
+from .util.service import handle_service_control
+
+
+@operation()
+def service(
+    service: str,
+    running: bool = True,
+    restarted: bool = False,
+    reloaded: bool = False,
+    command: Optional[str] = None,
+    enabled: Optional[bool] = None,
+    managed: bool = True,
+    svdir: str = "/var/service",
+    sourcedir: str = "/etc/sv",
+):
+    """
+    Manage the state of runit services.
+
+    + service: name of the service to manage
+    + running: whether the service should be running
+    + restarted: whether the service should be restarted
+    + reloaded: whether the service should be reloaded
+    + command: custom command to pass like: ``sv <command> <service>``
+    + enabled: whether this service should be enabled/disabled on boot
+    + managed: whether runit should manage this service
+
+      For services to be controlled, they first need to be managed by runit by
+      adding a symlink to the service in ``SVDIR``.
+      By setting ``managed=False`` the symlink will be removed.
+      Other options won't have any effect after that.
+      Although the ``<service>/down`` file can still be controlled with the
+      ``enabled`` option.
+
+    + svdir: alternative ``SVDIR``
+
+      An alternative ``SVDIR`` can be specified. This can be used for user services.
+
+    + sourcedir: where to search for available services
+
+      An alternative directory for available services can be specified.
+      Example: ``sourcedir=/etc/sv.local`` for services managed by the administrator.
+    """
+
+    was_managed = service in host.get_fact(RunitManaged, service=service, svdir=svdir)
+    was_auto = not host.get_fact(File, path="{0}/{1}/down".format(sourcedir, service))
+
+    # Disable autostart for previously unmanaged services.
+    #
+    # Where ``running=False`` is requested, this prevents one case of briefly
+    # starting and stopping the service.
+    if not was_managed and managed and was_auto:
+        yield from auto._inner(
+            service=service,
+            auto=False,
+            sourcedir=sourcedir,
+        )
+
+    yield from manage._inner(
+        service=service,
+        managed=managed,
+        svdir=svdir,
+        sourcedir=sourcedir,
+    )
+
+    # Service wasn't managed before, so wait for ``runsv`` to start.
+    # ``runsvdir`` will check at least every 5 seconds for new services.
+    # Wait for at most 10 seconds for the service to be managed, otherwise fail.
+    if not was_managed and managed:
+        yield from wait_runsv._inner(
+            service=service,
+            svdir=svdir,
+        )
+
+    if isinstance(enabled, bool):
+        yield from auto._inner(
+            service=service,
+            auto=enabled,
+            sourcedir=sourcedir,
+        )
+    else:
+        # restore previous state of ``<service>/down``
+        yield from auto._inner(
+            service=service,
+            auto=was_auto,
+            sourcedir=sourcedir,
+        )
+
+    # Services need to be managed by ``runit`` for the other options to make sense.
+    if not managed:
+        return
+
+    yield from handle_service_control(
+        host,
+        service,
+        host.get_fact(RunitStatus, service=service, svdir=svdir),
+        "SVDIR={0} sv {{1}} {{0}}".format(svdir),
+        running,
+        restarted,
+        reloaded,
+        command,
+    )
+
+
+@operation()
+def manage(
+    service: str,
+    managed: bool = True,
+    svdir: str = "/var/service",
+    sourcedir: str = "/etc/sv",
+):
+    """
+    Manage runit svdir links.
+
+    + service: name of the service to manage
+    + managed: whether the link should exist
+    + svdir: alternative ``SVDIR``
+    + sourcedir: where to search for available services
+    """
+
+    yield from link._inner(
+        path="{0}/{1}".format(svdir, service),
+        target="{0}/{1}".format(sourcedir, service),
+        present=managed,
+        create_remote_dir=False,
+    )
+
+
+@operation(is_idempotent=False)
+def wait_runsv(
+    service: str,
+    svdir: str = "/var/service",
+    timeout: int = 10,
+):
+    """
+    Wait for runsv for ``service`` to be available.
+
+    + service: name of the service to manage
+    + svdir: alternative ``SVDIR``
+    + timeout: time in seconds to wait
+    """
+
+    yield (
+        "export SVDIR={0}\n"
+        "for i in $(seq {1}); do\n"
+        "    sv status {2} > /dev/null && exit 0\n"
+        "    sleep 1;\n"
+        "done\n"
+        "exit 1"
+    ).format(svdir, timeout, service)
+
+
+@operation()
+def auto(
+    service: str,
+    auto: bool = True,
+    sourcedir: str = "/etc/sv",
+):
+    """
+    Start service automatically by managing the ``service/down`` file.
+
+    + service: name of the service to manage
+    + auto: whether the service should start automatically
+    + sourcedir: where to search for available services
+    """
+
+    yield from file._inner(
+        path="{0}/{1}/down".format(sourcedir, service),
+        present=not auto,
+        create_remote_dir=False,
+    )

pyinfra/operations/selinux.py

@@ -1,19 +1,36 @@
 """
 Provides operations to set SELinux file contexts, booleans and port types.
 """
+
+from __future__ import annotations
+
+from enum import Enum
+
 from pyinfra import host
-from pyinfra.api import QuoteString, StringCommand, operation
+from pyinfra.api import OperationValueError, QuoteString, StringCommand, operation
 from pyinfra.facts.selinux import FileContext, FileContextMapping, SEBoolean, SEPort, SEPorts
 from pyinfra.facts.server import Which
 
 
+class Boolean(Enum):
+    ON = "on"
+    OFF = "off"
+
+
+class Protocol(Enum):
+    UDP = "udp"
+    TCP = "tcp"
+    SCTP = "sctp"
+    DCCP = "dccp"
+
+
 @operation()
-def boolean(bool_name, value, persistent=False):
+def boolean(bool_name: str, value: Boolean, persistent=False):
     """
     Set the specified SELinux boolean to the desired state.
 
     + boolean: name of an SELinux boolean
-    + state: 'on' or 'off'
+    + value: desired state of the boolean
     + persistent: whether to write updated policy or not
 
     Note: This operation requires root privileges.
@@ -25,26 +42,31 @@ def boolean(bool_name, value, persistent=False):
         selinux.boolean(
             name='Allow Apache to connect to LDAP server',
             'httpd_can_network_connect',
-            'on',
+            Boolean.ON,
             persistent=True
         )
     """
-    _valid_states = ["on", "off"]
 
-    if value not in _valid_states:
-        raise ValueError(
-            f'\'value\' must be one of \'{",".join(_valid_states)}\' but found \'{value}\'',
-        )
+    value_str: str
+    if value in ["on", "off"]:  # compatibility with the old version
+        assert isinstance(value, str)
+        value_str = value
+    elif value is Boolean.ON:
+        value_str = "on"
+    elif value is Boolean.OFF:
+        value_str = "off"
+    else:
+        raise OperationValueError(f"Invalid value '{value}' for boolean operation")
 
-    if host.get_fact(SEBoolean, boolean=bool_name) != value:
+    if host.get_fact(SEBoolean, boolean=bool_name) != value_str:
         persist = "-P " if persistent else ""
-        yield StringCommand("setsebool", f"{persist}{bool_name}", value)
+        yield StringCommand("setsebool", f"{persist}{bool_name}", value_str)
     else:
-        host.noop(f"boolean '{bool_name}' already had the value '{value}'")
+        host.noop(f"boolean '{bool_name}' already had the value '{value_str}'")
 
 
 @operation()
-def file_context(path, se_type):
+def file_context(path: str, se_type: str):
     """
     Set the SELinux type for the specified path to the specified value.
 
@@ -70,7 +92,7 @@ def file_context(path, se_type):
 
 
 @operation()
-def file_context_mapping(target, se_type=None, present=True):
+def file_context_mapping(target: str, se_type: str | None = None, present=True):
     """
     Set the SELinux file context mapping for paths matching the target.
 
@@ -110,7 +132,7 @@ def file_context_mapping(target, se_type=None, present=True):
 
 
 @operation()
-def port(protocol, port_num, se_type=None, present=True):
+def port(protocol: Protocol | str, port_num: int, se_type: str | None = None, present=True):
     """
     Set the SELinux type for the specified protocol and port.
 
@@ -127,12 +149,16 @@ def port(protocol, port_num, se_type=None, present=True):
 
         selinux.port(
             name='Allow Apache to provide service on port 2222',
-            'tcp',
+            Protocol.TCP,
             2222,
             'http_port_t',
         )
     """
 
+    if protocol is Protocol:
+        assert isinstance(protocol, Protocol)
+        protocol = protocol.value
+
     if present and (se_type is None):
         raise ValueError("se_type must have a valid value if present is set")
 

pyinfra/operations/server.py

@@ -3,6 +3,8 @@ The server module takes care of os-level state. Targets POSIX compatibility, tes
 Linux/BSD.
 """
 
+from __future__ import annotations
+
 import shlex
 from io import StringIO
 from itertools import filterfalse, tee
@@ -18,6 +20,7 @@ from pyinfra.facts.files import Directory, FindInFile, Link
 from pyinfra.facts.server import (
     Crontab,
     Groups,
+    Home,
     Hostname,
     KernelModules,
     Locales,
@@ -38,6 +41,7 @@ from . import (
     openrc,
     pacman,
     pkg,
+    runit,
     systemd,
     sysvinit,
     upstart,
@@ -139,7 +143,7 @@ def wait(port: int):
 
 
 @operation(is_idempotent=False)
-def shell(commands):
+def shell(commands: str | list[str]):
     """
     Run raw shell code on server during a deploy. If the command would
     modify data that would be in a fact, the fact would not be updated
@@ -166,7 +170,7 @@ def shell(commands):
 
 
 @operation(is_idempotent=False)
-def script(src, args=()):
+def script(src: str, args=()):
     """
     Upload and execute a local script on the remote host.
 
@@ -199,7 +203,7 @@ def script(src, args=()):
 
 
 @operation(is_idempotent=False)
-def script_template(src, args=(), **data):
+def script_template(src: str, args=(), **data):
     """
     Generate, upload and execute a local script template on the remote host.
 
@@ -229,7 +233,7 @@ def script_template(src, args=(), **data):
 
 
 @operation()
-def modprobe(module, present=True, force=False):
+def modprobe(module: str, present=True, force=False):
     """
     Load/unload kernel modules.
 
@@ -281,11 +285,11 @@ def modprobe(module, present=True, force=False):
 
 @operation()
 def mount(
-    path,
+    path: str,
     mounted=True,
-    options=None,
-    device=None,
-    fs_type=None,
+    options: list[str] | None = None,
+    device: str | None = None,
+    fs_type: str | None = None,
     # TODO: do we want to manage fstab here?
     # update_fstab=False,
 ):
@@ -344,7 +348,7 @@ def mount(
 
 
 @operation()
-def hostname(hostname, hostname_file=None):
+def hostname(hostname: str, hostname_file: str | None = None):
     """
     Set the system hostname using ``hostnamectl`` or ``hostname`` on older systems.
 
@@ -402,8 +406,8 @@ def hostname(hostname, hostname_file=None):
 
 @operation()
 def sysctl(
-    key,
-    value,
+    key: str,
+    value: str | int | list[str | int],
     persist=False,
     persist_file="/etc/sysctl.conf",
 ):
@@ -449,12 +453,12 @@ def sysctl(
 
 @operation()
 def service(
-    service,
+    service: str,
     running=True,
     restarted=False,
     reloaded=False,
-    command=None,
-    enabled=None,
+    command: str | None = None,
+    enabled: bool | None = None,
 ):
     """
     Manage the state of services. This command checks for the presence of all the
@@ -489,6 +493,9 @@ def service(
     elif host.get_fact(Which, command="initctl"):
         service_operation = upstart.service
 
+    elif host.get_fact(Which, command="sv"):
+        service_operation = runit.service
+
     elif (
         host.get_fact(Which, command="service")
         or host.get_fact(Link, path="/etc/init.d")
@@ -518,7 +525,7 @@ def service(
 
 @operation()
 def packages(
-    packages,
+    packages: str | list[str],
     present=True,
 ):
     """
@@ -583,16 +590,16 @@ def packages(
 
 @operation()
 def crontab(
-    command,
+    command: str,
     present=True,
-    user=None,
-    cron_name=None,
+    user: str | None = None,
+    cron_name: str | None = None,
     minute="*",
     hour="*",
     month="*",
     day_of_week="*",
     day_of_month="*",
-    special_time=None,
+    special_time: str | None = None,
     interpolate_variables=False,
 ):
     """
@@ -663,7 +670,7 @@ def crontab(
 
     exists = existing_crontab is not None
 
-    edit_commands = []
+    edit_commands: list[str | StringCommand] = []
     temp_filename = host.get_temp_filename()
 
     if special_time:
@@ -758,7 +765,7 @@ def crontab(
 
 
 @operation()
-def group(group, present=True, system=False, gid=None):
+def group(group: str, present=True, system=False, gid: int | str | None = None):
     """
     Add/remove system groups.
 
@@ -827,12 +834,12 @@ def group(group, present=True, system=False, gid=None):
 
 @operation()
 def user_authorized_keys(
-    user,
-    public_keys,
-    group=None,
+    user: str,
+    public_keys: str | list[str],
+    group: str | None = None,
     delete_keys=False,
-    authorized_key_directory=None,
-    authorized_key_filename=None,
+    authorized_key_directory: str | None = None,
+    authorized_key_filename: str | None = None,
 ):
     """
     Manage `authorized_keys` of system users.
@@ -858,7 +865,9 @@ def user_authorized_keys(
     """
 
     if not authorized_key_directory:
-        authorized_key_directory = f"/home/{user}/.ssh/"
+        home = host.get_fact(Home, user=user)
+        authorized_key_directory = f"{home}/.ssh"
+
     if not authorized_key_filename:
         authorized_key_filename = "authorized_keys"
 
@@ -923,22 +932,22 @@ def user_authorized_keys(
 
 @operation()
 def user(
-    user,
+    user: str,
     present=True,
-    home=None,
-    shell=None,
-    group=None,
-    groups=None,
-    public_keys=None,
+    home: str | None = None,
+    shell: str | None = None,
+    group: str | None = None,
+    groups: list[str] | None = None,
+    public_keys: str | list[str] | None = None,
     delete_keys=False,
     ensure_home=True,
     create_home=False,
     system=False,
-    uid=None,
-    comment=None,
+    uid: int | None = None,
+    comment: str | None = None,
     add_deploy_dir=True,
     unique=True,
-    password=None,
+    password: str | None = None,
 ):
     """
     Add/remove/update system users & their ssh `authorized_keys`.
@@ -1124,7 +1133,7 @@ def user(
                 existing_user["password"] = password
 
     # Ensure home directory ownership
-    if ensure_home:
+    if ensure_home and home:
         yield from files.directory._inner(
             path=home,
             user=user,
@@ -1147,7 +1156,7 @@ def user(
 
 @operation()
 def locale(
-    locale,
+    locale: str,
     present=True,
 ):
     """
@@ -1217,10 +1226,10 @@ def locale(
 
 @operation()
 def security_limit(
-    domain,
-    limit_type,
-    item,
-    value,
+    domain: str,
+    limit_type: str,
+    item: str,
+    value: int,
 ):
     """
     Edit /etc/security/limits.conf configuration.
@@ -1239,7 +1248,7 @@ def security_limit(
             domain='*',
             limit_type='soft',
             item='nofile',
-            value='1024',
+            value=1024,
         )
     """
 

pyinfra/operations/snap.py

@@ -2,6 +2,8 @@
 Manage snap packages. See https://snapcraft.io/
 """
 
+from __future__ import annotations
+
 from pyinfra import host
 from pyinfra.api import operation
 from pyinfra.facts.snap import SnapPackage, SnapPackages
@@ -9,7 +11,7 @@ from pyinfra.facts.snap import SnapPackage, SnapPackages
 
 @operation()
 def package(
-    packages=None,
+    packages: str | list[str] | None = None,
     channel="latest/stable",
     classic=False,
     present=True,

pyinfra/operations/ssh.py

@@ -4,6 +4,8 @@ Execute commands and up/download files *from* the remote host.
 Eg: ``pyinfra -> inventory-host.net <-> another-host.net``
 """
 
+from __future__ import annotations
+
 import shlex
 
 from pyinfra import host
@@ -15,7 +17,7 @@ from . import files
 
 
 @operation()
-def keyscan(hostname, force=False, port=22):
+def keyscan(hostname: str, force=False, port=22):
     """
     Check/add hosts to the ``~/.ssh/known_hosts`` file.
 
@@ -63,7 +65,7 @@ def keyscan(hostname, force=False, port=22):
 
 
 @operation(is_idempotent=False)
-def command(hostname, command, user=None, port=22):
+def command(hostname: str, command: str, user: str | None = None, port=22):
     """
     Execute commands on other servers over SSH.
 
@@ -95,11 +97,11 @@ def command(hostname, command, user=None, port=22):
 
 @operation(is_idempotent=False)
 def upload(
-    hostname,
-    filename,
-    remote_filename=None,
+    hostname: str,
+    filename: str,
+    remote_filename: str | None = None,
     port=22,
-    user=None,
+    user: str | None = None,
     use_remote_sudo=False,
     ssh_keyscan=False,
 ):
@@ -159,12 +161,12 @@ def upload(
 
 @operation()
 def download(
-    hostname,
-    filename,
-    local_filename=None,
+    hostname: str,
+    filename: str,
+    local_filename: str | None = None,
     force=False,
     port=22,
-    user=None,
+    user: str | None = None,
     ssh_keyscan=False,
 ):
     """