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

pyinfra/api/__init__.py

@@ -11,6 +11,9 @@ from .config import Config  # noqa: F401 # pragma: no cover
 from .deploy import deploy  # noqa: F401 # pragma: no cover
 from .exceptions import DeployError  # noqa: F401 # pragma: no cover
 from .exceptions import (  # noqa: F401
+    FactError,
+    FactTypeError,
+    FactValueError,
     InventoryError,
     OperationError,
     OperationTypeError,

pyinfra/api/arguments.py

@@ -1,31 +1,42 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, Optional, Tuple, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    get_type_hints,
+)
 
 from typing_extensions import TypedDict
 
 from pyinfra import context
+from pyinfra.api.exceptions import ArgumentTypeError
 from pyinfra.api.state import State
-
-from .util import memoize
+from pyinfra.api.util import raise_if_bad_type
 
 if TYPE_CHECKING:
     from pyinfra.api.config import Config
     from pyinfra.api.host import Host
 
-
 T = TypeVar("T")
-
-
 default_sentinel = object()
 
 
-class ArgumentMeta:
+class ArgumentMeta(Generic[T]):
     description: str
     default: Callable[["Config"], T]
-    handler: Callable[["Config", T], T]
+    handler: Optional[Callable[["Config", T], T]]
 
-    def __init__(self, description, default, handler=default_sentinel) -> None:
+    def __init__(self, description, default, handler=None) -> None:
         self.description = description
         self.default = default
         self.handler = handler
@@ -36,12 +47,14 @@ class ArgumentMeta:
 # API to read/write external systems.
 
 
+# Note: ConnectorArguments is specifically not total as it's used to type many
+# functions via Unpack and we don't want to specify every kwarg.
 class ConnectorArguments(TypedDict, total=False):
     # Auth arguments
     _sudo: bool
     _sudo_user: str
     _use_sudo_login: bool
-    _use_sudo_password: bool
+    _sudo_password: str
     _preserve_sudo_env: bool
     _su_user: str
     _use_su_login: bool
@@ -59,7 +72,7 @@ class ConnectorArguments(TypedDict, total=False):
     _success_exit_codes: Iterable[int]
     _timeout: int
     _get_pty: bool
-    _stdin: Union[str, list, tuple]
+    _stdin: Union[str, Iterable[str]]
 
 
 def generate_env(config: "Config", value: dict) -> dict:
@@ -70,82 +83,82 @@ def generate_env(config: "Config", value: dict) -> dict:
 
 auth_argument_meta: dict[str, ArgumentMeta] = {
     "_sudo": ArgumentMeta(
-        "Execute/apply any changes with ``sudo``.",
+        "Execute/apply any changes with sudo.",
         default=lambda config: config.SUDO,
     ),
     "_sudo_user": ArgumentMeta(
-        "Execute/apply any changes with ``sudo`` as a non-root user.",
+        "Execute/apply any changes with sudo as a non-root user.",
         default=lambda config: config.SUDO_USER,
     ),
     "_use_sudo_login": ArgumentMeta(
-        "Execute ``sudo`` with a login shell.",
+        "Execute sudo with a login shell.",
         default=lambda config: config.USE_SUDO_LOGIN,
     ),
-    "_use_sudo_password": ArgumentMeta(
-        "Whether to use a password with ``sudo`` (will ask).",
-        default=lambda config: config.USE_SUDO_PASSWORD,
+    "_sudo_password": ArgumentMeta(
+        "Password to sudo with. If needed and not specified pyinfra will prompt for it.",
+        default=lambda config: config.SUDO_PASSWORD,
     ),
     "_preserve_sudo_env": ArgumentMeta(
-        "Preserve the shell environment when using ``sudo``.",
+        "Preserve the shell environment of the connecting user when using sudo.",
         default=lambda config: config.PRESERVE_SUDO_ENV,
     ),
     "_su_user": ArgumentMeta(
-        "Execute/apply any changes with this user using ``su``.",
+        "Execute/apply any changes with this user using su.",
         default=lambda config: config.SU_USER,
     ),
     "_use_su_login": ArgumentMeta(
-        "Execute ``su`` with a login shell.",
+        "Execute su with a login shell.",
         default=lambda config: config.USE_SU_LOGIN,
     ),
     "_preserve_su_env": ArgumentMeta(
-        "Preserve the shell environment when using ``su``.",
+        "Preserve the shell environment of the connecting user when using su.",
         default=lambda config: config.PRESERVE_SU_ENV,
     ),
     "_su_shell": ArgumentMeta(
-        "Use this shell (instead of user login shell) when using ``su``). "
-        "Only available under Linux, for use when using `su` with a user that "
-        "has nologin/similar as their login shell.",
+        "Use this shell (instead of user login shell) when using ``_su``). "
+        + "Only available under Linux, for use when using `su` with a user that "
+        + "has nologin/similar as their login shell.",
         default=lambda config: config.SU_SHELL,
     ),
     "_doas": ArgumentMeta(
-        "Execute/apply any changes with ``doas``.",
+        "Execute/apply any changes with doas.",
         default=lambda config: config.DOAS,
     ),
     "_doas_user": ArgumentMeta(
-        "Execute/apply any changes with ``doas`` as a non-root user.",
+        "Execute/apply any changes with doas as a non-root user.",
         default=lambda config: config.DOAS_USER,
     ),
 }
 
 shell_argument_meta: dict[str, ArgumentMeta] = {
     "_shell_executable": ArgumentMeta(
-        "The shell to use. Defaults to ``sh`` (Unix) or ``cmd`` (Windows).",
+        "The shell executable to use for executing commands.",
         default=lambda config: config.SHELL,
     ),
     "_chdir": ArgumentMeta(
         "Directory to switch to before executing the command.",
-        default=lambda config: "",
+        default=lambda _: "",
     ),
     "_env": ArgumentMeta(
         "Dictionary of environment variables to set.",
-        default=lambda config: {},
+        default=lambda _: {},
         handler=generate_env,
     ),
     "_success_exit_codes": ArgumentMeta(
         "List of exit codes to consider a success.",
-        default=lambda config: [0],
+        default=lambda _: [0],
     ),
     "_timeout": ArgumentMeta(
         "Timeout for *each* command executed during the operation.",
-        default=lambda config: None,
+        default=lambda _: None,
     ),
     "_get_pty": ArgumentMeta(
         "Whether to get a pseudoTTY when executing any commands.",
-        default=lambda config: False,
+        default=lambda _: False,
     ),
     "_stdin": ArgumentMeta(
         "String or buffer to send to the stdin of any commands.",
-        default=lambda config: None,
+        default=lambda _: None,
     ),
 }
 
@@ -154,21 +167,18 @@ shell_argument_meta: dict[str, ArgumentMeta] = {
 # These provide/extend additional operation metadata
 
 
-class MetaArguments(TypedDict, total=False):
+class MetaArguments(TypedDict):
     name: str
     _ignore_errors: bool
     _continue_on_error: bool
-    _precondition: str
-    _postcondition: str
-    _on_success: Callable[[State, "Host", str], None]
-    _on_error: Callable[[State, "Host", str], None]
+    _if: Union[List[Callable[[], bool]], Callable[[], bool], None]
 
 
 meta_argument_meta: dict[str, ArgumentMeta] = {
     # NOTE: name is the only non-_-prefixed argument
     "name": ArgumentMeta(
         "Name of the operation.",
-        default=lambda config: None,
+        default=lambda _: None,
     ),
     "_ignore_errors": ArgumentMeta(
         "Ignore errors when executing the operation.",
@@ -179,24 +189,11 @@ meta_argument_meta: dict[str, ArgumentMeta] = {
             "Continue executing operation commands after error. "
             "Only applies when ``_ignore_errors`` is true."
         ),
-        default=lambda config: False,
-    ),
-    "_precondition": ArgumentMeta(
-        "Command to execute & check before the operation commands begin.",
-        default=lambda config: None,
+        default=lambda _: False,
     ),
-    "_postcondition": ArgumentMeta(
-        "Command to execute & check after the operation commands complete.",
-        default=lambda config: None,
-    ),
-    # Lambda on the next two are to workaround a circular import
-    "_on_success": ArgumentMeta(
-        "Callback function to execute on success.",
-        default=lambda config: None,
-    ),
-    "_on_error": ArgumentMeta(
-        "Callback function to execute on error.",
-        default=lambda config: None,
+    "_if": ArgumentMeta(
+        "Only run this operation if these functions return True",
+        default=lambda _: [],
     ),
 }
 
@@ -206,7 +203,7 @@ meta_argument_meta: dict[str, ArgumentMeta] = {
 # over every target host for the same operation.
 
 
-class ExecutionArguments(TypedDict, total=False):
+class ExecutionArguments(TypedDict):
     _parallel: int
     _run_once: bool
     _serial: bool
@@ -219,11 +216,11 @@ execution_argument_meta: dict[str, ArgumentMeta] = {
     ),
     "_run_once": ArgumentMeta(
         "Only execute this operation once, on the first host to see it.",
-        default=lambda config: False,
+        default=lambda _: False,
     ),
     "_serial": ArgumentMeta(
         "Run this operation host by host, rather than in parallel.",
-        default=lambda config: False,
+        default=lambda _: False,
     ),
 }
 
@@ -232,6 +229,11 @@ class AllArguments(ConnectorArguments, MetaArguments, ExecutionArguments):
     pass
 
 
+def all_global_arguments() -> List[tuple[str, Type]]:
+    """Return all global arguments and their types."""
+    return list(get_type_hints(AllArguments).items())
+
+
 all_argument_meta: dict[str, ArgumentMeta] = {
     **auth_argument_meta,
     **shell_argument_meta,
@@ -239,29 +241,47 @@ all_argument_meta: dict[str, ArgumentMeta] = {
     **execution_argument_meta,
 }
 
-
-OPERATION_KWARG_DOC: list[tuple[str, Optional[str], type, dict[str, ArgumentMeta]]] = [
-    ("Privilege & user escalation", None, ConnectorArguments, auth_argument_meta),
-    ("Shell control & features", None, ConnectorArguments, shell_argument_meta),
-    ("Operation meta & callbacks", "Not available in facts.", MetaArguments, meta_argument_meta),
-    (
-        "Execution strategy",
-        "Not available in facts, value must be the same for all hosts.",
-        ExecutionArguments,
-        execution_argument_meta,
+EXECUTION_KWARG_KEYS = list(ExecutionArguments.__annotations__.keys())
+CONNECTOR_ARGUMENT_KEYS = list(ConnectorArguments.__annotations__.keys())
+
+__argument_docs__ = {
+    "Privilege & user escalation": (
+        auth_argument_meta,
+        """
+        .. code:: python
+
+            # Execute a command with sudo
+            server.user(
+                name="Create pyinfra user using sudo",
+                user="pyinfra",
+                _sudo=True,
+            )
+
+            # Execute a command with a specific sudo password
+            server.user(
+                name="Create pyinfra user using sudo",
+                user="pyinfra",
+                _sudo=True,
+                _sudo_password="my-secret-password",
+            )
+        """,
     ),
-]
-
-
-# Called ONCE to dedupe args that must be same for all calls of an op
-@memoize
-def get_execution_kwarg_keys() -> list[str]:
-    return list(ExecutionArguments.__annotations__.keys())
-
-
-@memoize
-def get_connector_argument_keys() -> list[str]:
-    return list(ConnectorArguments.__annotations__.keys())
+    "Shell control & features": (
+        shell_argument_meta,
+        """
+        .. code:: python
+
+            # Execute from a specific directory
+            server.shell(
+                name="Bootstrap nginx params",
+                commands=["openssl dhparam -out dhparam.pem 4096"],
+                _chdir="/etc/ssl/certs",
+            )
+        """,
+    ),
+    "Operation meta & callbacks": (meta_argument_meta, ""),
+    "Execution strategy": (execution_argument_meta, ""),
+}
 
 
 def pop_global_arguments(
@@ -269,7 +289,7 @@ def pop_global_arguments(
     state: Optional["State"] = None,
     host: Optional["Host"] = None,
     keys_to_check=None,
-) -> Tuple[AllArguments, list[str]]:
+) -> tuple[AllArguments, list[str]]:
     """
     Pop and return operation global keyword arguments, in preferred order:
 
@@ -277,16 +297,6 @@ def pop_global_arguments(
     + From any current @deploy context (deploy kwargs)
     + From the host data variables
     + From the config variables
-
-    Note this function is only called directly in the @operation & @deploy decorator
-    wrappers which the user should pass global arguments prefixed "_". This is to
-    avoid any clashes with operation and deploy functions both internal and third
-    party.
-
-    This is a bit strange because internally pyinfra uses non-_-prefixed arguments,
-    and this function is responsible for the translation between the two.
-
-    TODO: is this weird-ness acceptable? Is it worth updating internal use to _prefix?
     """
 
     state = state or context.state
@@ -296,12 +306,12 @@ def pop_global_arguments(
     if context.ctx_config.isset():
         config = context.config
 
-    meta_kwargs = host.current_deploy_kwargs or {}
+    meta_kwargs: dict[str, Any] = host.current_deploy_kwargs or {}  # type: ignore[assignment]
 
-    arguments: AllArguments = {}
+    arguments: dict[str, Any] = {}
     found_keys: list[str] = []
 
-    for key, type_ in AllArguments.__annotations__.items():
+    for key, type_ in all_global_arguments():
         if keys_to_check and key not in keys_to_check:
             continue
 
@@ -319,9 +329,17 @@ def pop_global_arguments(
         else:
             value = meta_kwargs.get(key, default)
 
-        if handler is not default_sentinel:
+        if handler:
             value = handler(config, value)
 
+        if value != default:
+            raise_if_bad_type(
+                value,
+                type_,
+                ArgumentTypeError,
+                f"Invalid argument `{key}`:",
+            )
+
         # TODO: why is type failing here?
         arguments[key] = value  # type: ignore
-    return arguments, found_keys
+    return cast(AllArguments, arguments), found_keys

pyinfra/api/arguments_typed.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Generator,
+    Generic,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Union,
+)
+
+from typing_extensions import ParamSpec, Protocol
+
+if TYPE_CHECKING:
+    from pyinfra.api.operation import OperationMeta
+
+P = ParamSpec("P")
+
+
+# Unfortunately we have to re-type out all of the global arguments here because
+# Python typing doesn't (yet) support merging kwargs. This acts as the operation
+# decorator output function which merges actual operation args/kwargs in paramspec
+# with the global arguments available in all operations.
+# The nature of "global arguments" is somewhat opposed to static typing as it's
+# indirect and somewhat magic, but we are where we are.
+class PyinfraOperation(Generic[P], Protocol):
+    _inner: Callable[P, Generator]
+
+    def __call__(
+        self,
+        #
+        # op args
+        # needs to be first
+        #
+        *args: P.args,
+        #
+        # ConnectorArguments
+        #
+        # Auth
+        _sudo: bool = False,
+        _sudo_user: Optional[str] = None,
+        _use_sudo_login: bool = False,
+        _sudo_password: Optional[str] = None,
+        _preserve_sudo_env: bool = False,
+        _su_user: Optional[str] = None,
+        _use_su_login: bool = False,
+        _preserve_su_env: bool = False,
+        _su_shell: Optional[str] = None,
+        _doas: bool = False,
+        _doas_user: Optional[str] = None,
+        # Shell arguments
+        _shell_executable: Optional[str] = None,
+        _chdir: Optional[str] = None,
+        _env: Optional[Mapping[str, str]] = None,
+        # Connector control
+        _success_exit_codes: Iterable[int] = (0,),
+        _timeout: Optional[int] = None,
+        _get_pty: bool = False,
+        _stdin: Union[None, str, list[str], tuple[str, ...]] = None,
+        #
+        # MetaArguments
+        #
+        name: Optional[str] = None,
+        _ignore_errors: bool = False,
+        _continue_on_error: bool = False,
+        _if: Union[List[Callable[[], bool]], Callable[[], bool], None] = None,
+        #
+        # ExecutionArguments
+        #
+        _parallel: Optional[int] = None,
+        _run_once: bool = False,
+        _serial: bool = False,
+        #
+        # op kwargs
+        #
+        **kwargs: P.kwargs,
+    ) -> "OperationMeta": ...

pyinfra/api/command.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import shlex
 from inspect import getfullargspec
 from string import Formatter
-from typing import TYPE_CHECKING, Callable, Union
+from typing import IO, TYPE_CHECKING, Callable, Union
 
 import gevent
 from typing_extensions import Unpack
@@ -143,7 +145,7 @@ class StringCommand(PyinfraCommand):
 class FileUploadCommand(PyinfraCommand):
     def __init__(
         self,
-        src: str,
+        src: str | IO,
         dest: str,
         remote_temp_filename=None,
         **kwargs: Unpack[ConnectorArguments],
@@ -173,7 +175,7 @@ class FileDownloadCommand(PyinfraCommand):
     def __init__(
         self,
         src: str,
-        dest: str,
+        dest: str | IO,
         remote_temp_filename=None,
         **kwargs: Unpack[ConnectorArguments],
     ):