Type stub generator for operations.

pyinfra/api/arguments.py

@@ -1,39 +1,51 @@
+from typing import Callable, Iterable, Mapping, Union
+
+import pyinfra
 from pyinfra import context, logger
+from pyinfra.api.state import State
 
 from .util import get_call_location, memoize
 
 auth_kwargs = {
     "_sudo": {
         "description": "Execute/apply any changes with ``sudo``.",
         "default": lambda config: config.SUDO,
+        "type": bool,
     },
     "_sudo_user": {
         "description": "Execute/apply any changes with ``sudo`` as a non-root user.",
         "default": lambda config: config.SUDO_USER,
+        "type": bool,
     },
     "_use_sudo_login": {
         "description": "Execute ``sudo`` with a login shell.",
         "default": lambda config: config.USE_SUDO_LOGIN,
+        "type": bool,
     },
     "_use_sudo_password": {
         "description": "Whether to use a password with ``sudo`` (will ask).",
         "default": lambda config: config.USE_SUDO_PASSWORD,
+        "type": bool,
     },
     "_preserve_sudo_env": {
         "description": "Preserve the shell environment when using ``sudo``.",
         "default": lambda config: config.PRESERVE_SUDO_ENV,
+        "type": bool,
     },
     "_su_user": {
-        "description": "Execute/apply any changes with ``su``.",
+        "description": "Execute/apply any changes with this user using ``su``.",
         "default": lambda config: config.SU_USER,
+        "type": str,
     },
     "_use_su_login": {
         "description": "Execute ``su`` with a login shell.",
         "default": lambda config: config.USE_SU_LOGIN,
+        "type": bool,
     },
     "_preserve_su_env": {
         "description": "Preserve the shell environment when using ``su``.",
         "default": lambda config: config.PRESERVE_SU_ENV,
+        "type": bool,
     },
     "_su_shell": {
         "description": (
@@ -42,14 +54,17 @@
             "has nologin/similar as their login shell."
         ),
         "default": lambda config: config.SU_SHELL,
+        "type": str,
     },
     "_doas": {
         "description": "Execute/apply any changes with ``doas``.",
         "defailt": lambda config: config.DOAS,
+        "type": bool,
     },
     "_doas_user": {
         "description": "Execute/apply any changes with ``doas`` as a non-root user.",
         "default": lambda config: config.DOAS_USER,
+        "type": str,
     },
 }
 
@@ -71,51 +86,82 @@ def generate_env(config, value):
     "_shell_executable": {
         "description": "The shell to use. Defaults to ``sh`` (Unix) or ``cmd`` (Windows).",
         "default": lambda config: config.SHELL,
+        "type": str,
     },
     "_chdir": {
         "description": "Directory to switch to before executing the command.",
+        "type": str,
     },
     "_env": {
         "description": "Dictionary of environment variables to set.",
         "handler": generate_env,
+        "type": Mapping[str, str],
     },
     "_success_exit_codes": {
         "description": "List of exit codes to consider a success.",
         "default": lambda config: [0],
+        "type": Iterable[int],
+    },
+    "_timeout": {
+        "description": "Timeout for *each* command executed during the operation.",
+        "type": int,
+    },
+    "_get_pty": {
+        "description": "Whether to get a pseudoTTY when executing any commands.",
+        "type": bool,
+    },
+    "_stdin": {
+        "description": "String or buffer to send to the stdin of any commands.",
+        "type": Union[str, list, tuple],
     },
-    "_timeout": "Timeout for *each* command executed during the operation.",
-    "_get_pty": "Whether to get a pseudoTTY when executing any commands.",
-    "_stdin": "String or buffer to send to the stdin of any commands.",
 }
 
 meta_kwargs = {
     # NOTE: name is the only non-_-prefixed argument
     "name": {
         "description": "Name of the operation.",
+        "type": str,
     },
     "_ignore_errors": {
         "description": "Ignore errors when executing the operation.",
         "default": lambda config: config.IGNORE_ERRORS,
+        "type": bool,
+    },
+    "_precondition": {
+        "description": "Command to execute & check before the operation commands begin.",
+        "type": str,
+    },
+    "_postcondition": {
+        "description": "Command to execute & check after the operation commands complete.",
+        "type": str,
+    },
+    # Lambda on the next two are to workaround a circular import
+    "_on_success": {
+        "description": "Callback function to execute on success.",
+        "type": lambda: Callable[[State, pyinfra.api.Host, str], None],
+    },
+    "_on_error": {
+        "description": "Callback function to execute on error.",
+        "type": lambda: Callable[[State, pyinfra.api.Host, str], None],
     },
-    "_precondition": "Command to execute & check before the operation commands begin.",
-    "_postcondition": "Command to execute & check after the operation commands complete.",
-    "_on_success": "Callback function to execute on success.",
-    "_on_error": "Callback function to execute on error.",
 }
 
 # Execution kwargs are global - ie must be identical for every host
 execution_kwargs = {
     "_parallel": {
         "description": "Run this operation in batches of hosts.",
         "default": lambda config: config.PARALLEL,
+        "type": int,
     },
     "_run_once": {
         "description": "Only execute this operation once, on the first host to see it.",
         "default": lambda config: False,
+        "type": bool,
     },
     "_serial": {
         "description": "Run this operation host by host, rather than in parallel.",
         "default": lambda config: False,
+        "type": bool,
     },
 }
 

pyinfra/operations/apk.py

@@ -10,7 +10,7 @@
 
 
 @operation(is_idempotent=False)
-def upgrade(available=False):
+def upgrade(available: bool = False):
     """
     Upgrades all apk packages.
 

pyinfra/operations/apk.pyi

@@ -0,0 +1,108 @@
+import typing
+import pyinfra
+
+def upgrade(
+    available: bool = False,
+    _sudo: typing.Optional[bool] = None,
+    _sudo_user: typing.Optional[bool] = None,
+    _use_sudo_login: typing.Optional[bool] = None,
+    _use_sudo_password: typing.Optional[bool] = None,
+    _preserve_sudo_env: typing.Optional[bool] = None,
+    _su_user: typing.Optional[str] = None,
+    _use_su_login: typing.Optional[bool] = None,
+    _preserve_su_env: typing.Optional[bool] = None,
+    _su_shell: typing.Optional[str] = None,
+    _doas: typing.Optional[bool] = None,
+    _doas_user: typing.Optional[str] = None,
+    _shell_executable: typing.Optional[str] = None,
+    _chdir: typing.Optional[str] = None,
+    _env: typing.Optional[typing.Mapping[str, str]] = None,
+    _success_exit_codes: typing.Optional[typing.Iterable[int]] = None,
+    _timeout: typing.Optional[int] = None,
+    _get_pty: typing.Optional[bool] = None,
+    _stdin: typing.Optional[typing.Union[str, list, tuple]] = None,
+    name: typing.Optional[str] = None,
+    _ignore_errors: typing.Optional[bool] = None,
+    _precondition: typing.Optional[str] = None,
+    _postcondition: typing.Optional[str] = None,
+    _on_success: typing.Optional[
+        typing.Callable[[pyinfra.api.state.State, pyinfra.api.host.Host, str], None]
+    ] = None,
+    _on_error: typing.Optional[
+        typing.Callable[[pyinfra.api.state.State, pyinfra.api.host.Host, str], None]
+    ] = None,
+    _parallel: typing.Optional[int] = None,
+    _run_once: typing.Optional[bool] = None,
+    _serial: typing.Optional[bool] = None,
+): ...
+def update(
+    _sudo: typing.Optional[bool] = None,
+    _sudo_user: typing.Optional[bool] = None,
+    _use_sudo_login: typing.Optional[bool] = None,
+    _use_sudo_password: typing.Optional[bool] = None,
+    _preserve_sudo_env: typing.Optional[bool] = None,
+    _su_user: typing.Optional[str] = None,
+    _use_su_login: typing.Optional[bool] = None,
+    _preserve_su_env: typing.Optional[bool] = None,
+    _su_shell: typing.Optional[str] = None,
+    _doas: typing.Optional[bool] = None,
+    _doas_user: typing.Optional[str] = None,
+    _shell_executable: typing.Optional[str] = None,
+    _chdir: typing.Optional[str] = None,
+    _env: typing.Optional[typing.Mapping[str, str]] = None,
+    _success_exit_codes: typing.Optional[typing.Iterable[int]] = None,
+    _timeout: typing.Optional[int] = None,
+    _get_pty: typing.Optional[bool] = None,
+    _stdin: typing.Optional[typing.Union[str, list, tuple]] = None,
+    name: typing.Optional[str] = None,
+    _ignore_errors: typing.Optional[bool] = None,
+    _precondition: typing.Optional[str] = None,
+    _postcondition: typing.Optional[str] = None,
+    _on_success: typing.Optional[
+        typing.Callable[[pyinfra.api.state.State, pyinfra.api.host.Host, str], None]
+    ] = None,
+    _on_error: typing.Optional[
+        typing.Callable[[pyinfra.api.state.State, pyinfra.api.host.Host, str], None]
+    ] = None,
+    _parallel: typing.Optional[int] = None,
+    _run_once: typing.Optional[bool] = None,
+    _serial: typing.Optional[bool] = None,
+): ...
+def packages(
+    packages=None,
+    present=True,
+    latest=False,
+    update=False,
+    upgrade=False,
+    _sudo: typing.Optional[bool] = None,
+    _sudo_user: typing.Optional[bool] = None,
+    _use_sudo_login: typing.Optional[bool] = None,
+    _use_sudo_password: typing.Optional[bool] = None,
+    _preserve_sudo_env: typing.Optional[bool] = None,
+    _su_user: typing.Optional[str] = None,
+    _use_su_login: typing.Optional[bool] = None,
+    _preserve_su_env: typing.Optional[bool] = None,
+    _su_shell: typing.Optional[str] = None,
+    _doas: typing.Optional[bool] = None,
+    _doas_user: typing.Optional[str] = None,
+    _shell_executable: typing.Optional[str] = None,
+    _chdir: typing.Optional[str] = None,
+    _env: typing.Optional[typing.Mapping[str, str]] = None,
+    _success_exit_codes: typing.Optional[typing.Iterable[int]] = None,
+    _timeout: typing.Optional[int] = None,
+    _get_pty: typing.Optional[bool] = None,
+    _stdin: typing.Optional[typing.Union[str, list, tuple]] = None,
+    name: typing.Optional[str] = None,
+    _ignore_errors: typing.Optional[bool] = None,
+    _precondition: typing.Optional[str] = None,
+    _postcondition: typing.Optional[str] = None,
+    _on_success: typing.Optional[
+        typing.Callable[[pyinfra.api.state.State, pyinfra.api.host.Host, str], None]
+    ] = None,
+    _on_error: typing.Optional[
+        typing.Callable[[pyinfra.api.state.State, pyinfra.api.host.Host, str], None]
+    ] = None,
+    _parallel: typing.Optional[int] = None,
+    _run_once: typing.Optional[bool] = None,
+    _serial: typing.Optional[bool] = None,
+): ...