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

pyinfra/api/util.py

@@ -1,70 +1,99 @@
-import re
-import shlex
+from __future__ import annotations
 
+import hashlib
 from functools import wraps
-from hashlib import sha1
+from hashlib import md5, sha1, sha256
 from inspect import getframeinfo, stack
-from socket import (
-    error as socket_error,
-    timeout as timeout_error,
-)
-from types import GeneratorType
+from io import BytesIO, StringIO
+from os import getcwd, path, stat
+from socket import error as socket_error, timeout as timeout_error
+from typing import IO, TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
 
 import click
-
-from jinja2 import Template, TemplateSyntaxError, UndefinedError
+from jinja2 import Environment, FileSystemLoader, StrictUndefined
 from paramiko import SSHException
+from typeguard import TypeCheckError, check_type
 
+import pyinfra
 from pyinfra import logger
-from pyinfra.api import Config
 
-from .exceptions import PyinfraError
+if TYPE_CHECKING:
+    from pyinfra.api.host import Host
+    from pyinfra.api.state import State, StateOperationMeta
+    from pyinfra.connectors.util import CommandOutput
 
 # 64kb chunks
 BLOCKSIZE = 65536
 
 # Caches
-TEMPLATES = {}
-FILE_SHAS = {}
+TEMPLATES: Dict[Any, Any] = {}
+FILE_SHAS: Dict[Any, Any] = {}
 
+PYINFRA_INSTALL_DIR = path.normpath(path.join(path.dirname(__file__), ".."))
 
-def try_int(value):
-    try:
-        return int(value)
-    except (TypeError, ValueError):
-        return value
 
+def get_file_path(state: "State", filename: str):
+    if path.isabs(filename):
+        return filename
+
+    assert state.cwd is not None, "Cannot use `get_file_path` with no `state.cwd` set"
+    relative_to = state.cwd
 
-def ensure_host_list(hosts, inventory):
-    if hosts is None:
-        return hosts
+    if state.current_exec_filename and (filename.startswith("./") or filename.startswith(".\\")):
+        relative_to = path.dirname(state.current_exec_filename)
 
-    # If passed a string, treat as group name and get any hosts from inventory
-    if isinstance(hosts, str):
-        return inventory.get_group(hosts, [])
+    return path.join(relative_to, filename)
 
-    if not isinstance(hosts, (list, tuple)):
-        return [hosts]
 
-    return hosts
+def get_kwargs_str(kwargs: Dict[Any, Any]):
+    if not kwargs:
+        return ""
 
+    items = [
+        "{0}={1}".format(key, value)
+        for key, value in sorted(kwargs.items())
+        if key not in ("self", "state", "host")
+    ]
+    return ", ".join(items)
 
-def memoize(func):
+
+def try_int(value):
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return value
+
+
+def memoize(func: Callable[..., Any]):
     @wraps(func)
     def wrapper(*args, **kwargs):
-        key = '{0}{1}'.format(args, kwargs)
-        if key in wrapper.cache:
-            return wrapper.cache[key]
+        key = "{0}{1}".format(args, kwargs)
+        if key in wrapper.cache:  # type: ignore[attr-defined]
+            return wrapper.cache[key]  # type: ignore[attr-defined]
 
         value = func(*args, **kwargs)
-        wrapper.cache[key] = value
+        wrapper.cache[key] = value  # type: ignore[attr-defined]
         return value
 
-    wrapper.cache = {}
+    wrapper.cache = {}  # type: ignore[attr-defined]
     return wrapper
 
 
-def get_caller_frameinfo(frame_offset=0):
+def get_call_location(frame_offset: int = 1):
+    frame = get_caller_frameinfo(frame_offset=frame_offset)  # escape *this* function
+    relpath = frame.filename
+
+    try:
+        # On Windows if pyinfra is on a different drive to the filename here, this will
+        # error as there's no way to do relative paths between drives.
+        relpath = path.relpath(frame.filename)
+    except ValueError:
+        pass
+
+    return "line {0} in {1}".format(frame.lineno, relpath)
+
+
+def get_caller_frameinfo(frame_offset: int = 0):
     # Default frames to look at is 2; one for this function call itself
     # in util.py and one for the caller of this function within pyinfra
     # giving the external call frame (ie end user deploy code).
@@ -80,394 +109,277 @@ def get_caller_frameinfo(frame_offset=0):
     return info
 
 
-def extract_callable_datas(datas):
-    for data in datas:
-        # Support for dynamic data, ie @deploy wrapped data defaults where
-        # the data is stored on the state temporarily.
-        if callable(data):
-            data = data()
+def get_operation_order_from_stack(state: "State"):
+    stack_items = list(reversed(stack()))
 
-        yield data
+    i = 0
+    # Find the *first* occurrence of our deploy file in the reversed stack
+    if state.current_deploy_filename:
+        for i, stack_item in enumerate(stack_items):
+            frame = getframeinfo(stack_item[0])
+            if frame.filename == state.current_deploy_filename:
+                break
 
+    # Now generate a list of line numbers *following that file*
+    line_numbers = []
 
-class FallbackDict(object):
-    '''
-    Combines multiple AttrData's to search for attributes.
-    '''
+    if pyinfra.is_cli:
+        line_numbers.append(state.current_op_file_number)
 
-    override_datas = None
+    for stack_item in stack_items[i:]:
+        frame = getframeinfo(stack_item[0])
 
-    def __init__(self, *datas):
-        datas = list(datas)
+        if frame.filename.startswith(PYINFRA_INSTALL_DIR):
+            continue
 
-        # Inject an empty override data so we can assign during deploy
-        self.__dict__['override_datas'] = {}
-        datas.insert(0, self.override_datas)
+        line_numbers.append(frame.lineno)
 
-        self.__dict__['datas'] = tuple(datas)
+    del stack_items
 
-    def __getattr__(self, key):
-        for data in extract_callable_datas(self.datas):
-            if key in data:
-                return data[key]
+    return line_numbers
 
-    def __setattr__(self, key, value):
-        self.override_datas[key] = value
 
-    def __str__(self):
-        return str(self.datas)
+def get_template(filename_or_io: str | IO, jinja_env_kwargs: dict[str, Any] | None = None):
+    """
+    Gets a jinja2 ``Template`` object for the input filename or string, with caching
+    based on the filename of the template, or the SHA1 of the input string.
+    """
+    if jinja_env_kwargs is None:
+        jinja_env_kwargs = {}
+    file_data = get_file_io(filename_or_io, mode="r")
+    cache_key = file_data.cache_key
 
-    def dict(self):
-        out = {}
+    if cache_key and cache_key in TEMPLATES:
+        return TEMPLATES[cache_key]
 
-        # Copy and reverse data objects (such that the first items override
-        # the last, matching __getattr__ output).
-        datas = list(self.datas)
-        datas.reverse()
+    with file_data as file_io:
+        template_string = file_io.read()
 
-        for data in extract_callable_datas(datas):
-            out.update(data)
+    template = Environment(
+        undefined=StrictUndefined,
+        keep_trailing_newline=True,
+        loader=FileSystemLoader(getcwd()),
+        **jinja_env_kwargs,
+    ).from_string(template_string)
 
-        return out
+    if cache_key:
+        TEMPLATES[cache_key] = template
 
+    return template
 
-def pop_op_kwargs(state, kwargs):
-    '''
-    Pop and return operation global keyword arguments.
-    '''
 
-    meta_kwargs = state.deploy_kwargs or {}
+def sha1_hash(string: str) -> str:
+    """
+    Return the SHA1 of the input string.
+    """
 
-    def get_kwarg(key, default=None):
-        return kwargs.pop(key, meta_kwargs.get(key, default))
+    hasher = sha1()
+    hasher.update(string.encode("utf-8"))
+    return hasher.hexdigest()
 
-    # Get the env for this host: config env followed by command-level env
-    env = state.config.ENV.copy()
-    env.update(get_kwarg('env', {}))
 
-    hosts = get_kwarg('hosts')
-    hosts = ensure_host_list(hosts, inventory=state.inventory)
+def format_exception(e: Exception) -> str:
+    return f"{e.__class__.__name__}{e.args}"
 
-    # Filter out any hosts not in the meta kwargs (nested support)
-    if meta_kwargs.get('hosts') is not None:
-        hosts = [
-            host for host in hosts
-            if host in meta_kwargs['hosts']
-        ]
 
-    return {
-        # ENVars for commands in this operation
-        'env': env,
-        # Hosts to limit the op to
-        'hosts': hosts,
-        # When to limit the op (default always)
-        'when': get_kwarg('when', True),
-        # Locally & globally configurable
-        'shell_executable': get_kwarg('shell_executable', state.config.SHELL),
-        'sudo': get_kwarg('sudo', state.config.SUDO),
-        'sudo_user': get_kwarg('sudo_user', state.config.SUDO_USER),
-        'su_user': get_kwarg('su_user', state.config.SU_USER),
-        # Whether to preserve ENVars when sudoing (eg SSH forward agent socket)
-        'preserve_sudo_env': get_kwarg(
-            'preserve_sudo_env', state.config.PRESERVE_SUDO_ENV,
-        ),
-        # Ignore any errors during this operation
-        'ignore_errors': get_kwarg(
-            'ignore_errors', state.config.IGNORE_ERRORS,
-        ),
-        # Timeout on running the command
-        'timeout': get_kwarg('timeout'),
-        # Get a PTY before executing commands
-        'get_pty': get_kwarg('get_pty', False),
-        # Forces serial mode for this operation (--serial for one op)
-        'serial': get_kwarg('serial', False),
-        # Only runs this operation once
-        'run_once': get_kwarg('run_once', False),
-        # Execute in batches of X hosts rather than all at once
-        'parallel': get_kwarg('parallel'),
-        # Callbacks
-        'on_success': get_kwarg('on_success'),
-        'on_error': get_kwarg('on_error'),
-        # Operation hash
-        'op': get_kwarg('op'),
-    }
-
-
-def unroll_generators(generator):
-    '''
-    Take a generator and unroll any sub-generators recursively. This is
-    essentially a Python 2 way of doing `yield from` in Python 3 (given
-    iterating the entire thing).
-    '''
-
-    # Ensure we have a generator (prevents ccommands returning lists)
-    if not isinstance(generator, GeneratorType):
-        raise TypeError('{0} is not a generator'.format(generator))
-
-    items = []
-
-    for item in generator:
-        if isinstance(item, GeneratorType):
-            items.extend(unroll_generators(item))
+def print_host_combined_output(host: "Host", output: "CommandOutput") -> None:
+    for line in output:
+        if line.buffer_name == "stderr":
+            logger.error(f"{host.print_prefix}{click.style(line.line, 'red')}")
         else:
-            items.append(item)
-
-    return items
-
-
-def get_template(filename_or_string, is_string=False):
-    '''
-    Gets a jinja2 ``Template`` object for the input filename or string, with caching
-    based on the filename of the template, or the SHA1 of the input string.
-    '''
-
-    # Cache against string sha or just the filename
-    cache_key = sha1_hash(filename_or_string) if is_string else filename_or_string
-
-    if cache_key in TEMPLATES:
-        return TEMPLATES[cache_key]
-
-    if is_string:
-        # Set the input string as our template
-        template_string = filename_or_string
-
-    else:
-        # Load template data into memory
-        with open(filename_or_string, 'r') as file_io:
-            template_string = file_io.read()
-
-    TEMPLATES[cache_key] = Template(template_string, keep_trailing_newline=True)
-    return TEMPLATES[cache_key]
-
-
-def underscore(name):
-    '''
-    Transform CamelCase -> snake_case.
-    '''
-
-    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
-    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()
-
+            logger.error(f"{host.print_prefix}{line.line}")
+
+
+def log_operation_start(
+    op_meta: "StateOperationMeta", op_types: Optional[List] = None, prefix: str = "--> "
+) -> None:
+    op_types = op_types or []
+    if op_meta.global_arguments["_serial"]:
+        op_types.append("serial")
+    if op_meta.global_arguments["_run_once"]:
+        op_types.append("run once")
+
+    args = ""
+    if op_meta.args:
+        args = "({0})".format(", ".join(str(arg) for arg in op_meta.args))
+
+    logger.info(
+        "{0} {1} {2}".format(
+            click.style(
+                "{0}Starting{1}operation:".format(
+                    prefix,
+                    " {0} ".format(", ".join(op_types)) if op_types else " ",
+                ),
+                "blue",
+            ),
+            click.style(", ".join(op_meta.names), bold=True),
+            args,
+        ),
+    )
 
-def sha1_hash(string):
-    '''
-    Return the SHA1 of the input string.
-    '''
 
-    hasher = sha1()
-    hasher.update(string.encode())
-    return hasher.hexdigest()
+def log_error_or_warning(
+    host: "Host", ignore_errors: bool, description: str = "", continue_on_error: bool = False
+) -> None:
+    log_func = logger.error
+    log_color = "red"
+    log_text = "Error: " if description else "Error"
 
+    if ignore_errors:
+        log_func = logger.warning
+        log_color = "yellow"
+        log_text = (
+            "Error (ignored, execution continued)" if continue_on_error else "Error (ignored)"
+        )
+        if description:
+            log_text = f"{log_text}: "
 
-def format_exception(e):
-    return '{0}{1}'.format(e.__class__.__name__, e.args)
+    log_func(
+        "{0}{1}{2}".format(
+            host.print_prefix,
+            click.style(log_text, log_color),
+            description,
+        ),
+    )
 
 
-def log_host_command_error(host, e, timeout=0):
+def log_host_command_error(host: "Host", e: Exception, timeout: int | None = 0) -> None:
     if isinstance(e, timeout_error):
-        logger.error('{0}{1}'.format(
-            host.print_prefix,
-            click.style('Command timed out after {0}s'.format(
-                timeout,
-            ), 'red'),
-        ))
+        logger.error(
+            "{0}{1}".format(
+                host.print_prefix,
+                click.style(
+                    "Command timed out after {0}s".format(
+                        timeout,
+                    ),
+                    "red",
+                ),
+            ),
+        )
 
     elif isinstance(e, (socket_error, SSHException)):
-        logger.error('{0}{1}'.format(
-            host.print_prefix,
-            click.style('Command socket/SSH error: {0}'.format(
-                format_exception(e)), 'red',
+        logger.error(
+            "{0}{1}".format(
+                host.print_prefix,
+                click.style(
+                    "Command socket/SSH error: {0}".format(format_exception(e)),
+                    "red",
+                ),
             ),
-        ))
+        )
 
     elif isinstance(e, IOError):
-        logger.error('{0}{1}'.format(
-            host.print_prefix,
-            click.style('Command IO error: {0}'.format(
-                format_exception(e)), 'red',
+        logger.error(
+            "{0}{1}".format(
+                host.print_prefix,
+                click.style(
+                    "Command IO error: {0}".format(format_exception(e)),
+                    "red",
+                ),
             ),
-        ))
+        )
 
     # Still here? Re-raise!
     else:
         raise e
 
 
-def make_command(
-    command,
-    env=None,
-    su_user=None,
-    sudo=False,
-    sudo_user=None,
-    preserve_sudo_env=False,
-    shell_executable=Config.SHELL,
-):
-    '''
-    Builds a shell command with various kwargs.
-    '''
-
-    debug_meta = {}
-
-    for key, value in (
-        ('shell_executable', shell_executable),
-        ('sudo', sudo),
-        ('sudo_user', sudo_user),
-        ('su_user', su_user),
-        ('env', env),
-    ):
-        if value:
-            debug_meta[key] = value
-
-    logger.debug('Building command ({0}): {1}'.format(' '.join(
-        '{0}: {1}'.format(key, value)
-        for key, value in debug_meta.items()
-    ), command))
-
-    # Use env & build our actual command
-    if env:
-        env_string = ' '.join([
-            '{0}={1}'.format(key, value)
-            for key, value in env.items()
-        ])
-        command = 'export {0}; {1}'.format(env_string, command)
-
-    # Quote the command as a string
-    command = shlex.quote(command)
-
-    # Switch user with su
-    if su_user:
-        # note `which <shell>` usage here - su requires an absolute path
-        command = 'su {0} -s `which {1}` -c {2}'.format(
-            su_user, shell_executable, command,
-        )
-
-    # Otherwise just sh wrap the command
-    else:
-        command = '{0} -c {1}'.format(shell_executable, command)
-
-    # Use sudo (w/user?)
-    if sudo:
-        sudo_bits = ['sudo', '-H']
-
-        if preserve_sudo_env:
-            sudo_bits.append('-E')
-
-        if sudo_user:
-            sudo_bits.extend(('-u', sudo_user))
-
-        command = '{0} {1}'.format(' '.join(sudo_bits), command)
-
-    return command
-
-
-def get_arg_value(state, host, arg):
-    '''
-    Runs string arguments through the jinja2 templating system with a state and
-    host. Used to avoid string formatting in deploy operations which result in
-    one operation per host/variable. By parsing the commands after we generate
-    the ``op_hash``, multiple command variations can fall under one op.
-    '''
-
-    if isinstance(arg, str):
-        data = {
-            'host': host,
-            'inventory': state.inventory,
-        }
-
-        try:
-            return get_template(arg, is_string=True).render(data)
-        except (TemplateSyntaxError, UndefinedError) as e:
-            raise PyinfraError('Error in template string: {0}'.format(e))
-
-    elif isinstance(arg, list):
-        return [get_arg_value(state, host, value) for value in arg]
-
-    elif isinstance(arg, tuple):
-        return tuple(get_arg_value(state, host, value) for value in arg)
-
-    elif isinstance(arg, dict):
-        return {
-            key: get_arg_value(state, host, value)
-            for key, value in arg.items()
-        }
-
-    return arg
-
-
 def make_hash(obj):
-    '''
+    """
     Make a hash from an arbitrary nested dictionary, list, tuple or set, used to generate
     ID's for operations based on their name & arguments.
-    '''
+    """
 
     if isinstance(obj, (set, tuple, list)):
-        hash_string = ''.join([make_hash(e) for e in obj])
+        hash_string = "".join([make_hash(e) for e in obj])
 
     elif isinstance(obj, dict):
-        hash_string = ''.join(
-            ''.join((key, make_hash(value)))
-            for key, value in obj.items()
-        )
+        hash_string = "".join("".join((key, make_hash(value))) for key, value in obj.items())
 
     else:
         hash_string = (
+            # Capture integers first (as 1 == True)
+            "{0}".format(obj)
+            if isinstance(obj, int)
             # Constants - the values can change between hosts but we should still
             # group them under the same operation hash.
-            '_PYINFRA_CONSTANT' if obj in (True, False, None)
-            # Plain strings
-            else obj if isinstance(obj, str)
-            # Objects with __name__s
-            else obj.__name__ if hasattr(obj, '__name__')
-            # Objects with names
-            else obj.name if hasattr(obj, 'name')
-            # Repr anything else
-            else repr(obj)
+            else (
+                "_PYINFRA_CONSTANT"
+                if obj in (True, False, None)
+                # Plain strings
+                else (
+                    obj
+                    if isinstance(obj, str)
+                    # Objects with __name__s
+                    else (
+                        obj.__name__
+                        if hasattr(obj, "__name__")
+                        # Objects with names
+                        else (
+                            obj.name
+                            if hasattr(obj, "name")
+                            # Repr anything else
+                            else repr(obj)
+                        )
+                    )
+                )
+            )
         )
 
     return sha1_hash(hash_string)
 
 
-class get_file_io(object):
-    '''
+class get_file_io:
+    """
     Given either a filename or an existing IO object, this context processor
     will open and close filenames, and leave IO objects alone.
-    '''
+    """
+
+    filename_or_io: Union[str, IO[Any]]
+    mode: str
 
-    close = False
+    _close: bool = False
+    _file_io: IO[Any]
 
-    def __init__(self, filename_or_io):
+    def __init__(self, filename_or_io: str | IO, mode: str = "rb"):
         if not (
             # Check we can be read
-            hasattr(filename_or_io, 'read')
+            hasattr(filename_or_io, "read")
             # Or we're a filename
             or isinstance(filename_or_io, str)
         ):
-            raise TypeError('Invalid filename or IO object: {0}'.format(
-                filename_or_io,
-            ))
+            raise TypeError(
+                "Invalid filename or IO object: {0}".format(
+                    filename_or_io,
+                ),
+            )
+
+        # Convert any StringIO/BytesIO to the other to match the desired mode
+        if isinstance(filename_or_io, StringIO) and mode == "rb":
+            filename_or_io.seek(0)
+            filename_or_io = BytesIO(filename_or_io.read().encode())
+        if isinstance(filename_or_io, BytesIO) and mode == "r":
+            filename_or_io.seek(0)
+            filename_or_io = StringIO(filename_or_io.read().decode())
 
         self.filename_or_io = filename_or_io
+        self.mode = mode
 
     def __enter__(self):
-        # If we have a read attribute, just use the object as-is
-        if hasattr(self.filename_or_io, 'read'):
-            file_io = self.filename_or_io
-
-        # Otherwise, assume a filename and open it up
+        if isinstance(self.filename_or_io, str):
+            file_io = open(self.filename_or_io, self.mode)
+            self._file_io = file_io
+            self._close = True
         else:
-            file_io = open(self.filename_or_io, 'rb')
-
-            # Attach to self for closing on __exit__
-            self.file_io = file_io
-            self.close = True
+            file_io = self.filename_or_io
 
         # Ensure we're at the start of the file
         file_io.seek(0)
         return file_io
 
     def __exit__(self, type, value, traceback):
-        if self.close:
-            self.file_io.close()
+        if self._close:
+            self._file_io.close()
 
     @property
     def cache_key(self):
@@ -477,24 +389,37 @@ class get_file_io(object):
             return self.filename_or_io
 
 
-def get_file_sha1(filename_or_io):
-    '''
-    Calculates the SHA1 of a file or file object using a buffer to handle larger files.
-    '''
+def get_file_md5(filename_or_io: str | IO):
+    return _get_file_digest(filename_or_io, md5())
+
+
+def get_file_sha1(filename_or_io: str | IO):
+    return _get_file_digest(filename_or_io, sha1())
+
+
+def get_file_sha256(filename_or_io: str | IO):
+    return _get_file_digest(filename_or_io, sha256())
+
+
+def _get_file_digest(filename_or_io: str | IO, hasher: hashlib._Hash):
+    """
+    Calculates the hash of a file or file object using a buffer to handle larger files.
+    """
 
     file_data = get_file_io(filename_or_io)
     cache_key = file_data.cache_key
+    if cache_key:
+        cache_key = f"{cache_key}_{hasher.name}"
 
     if cache_key and cache_key in FILE_SHAS:
         return FILE_SHAS[cache_key]
 
     with file_data as file_io:
-        hasher = sha1()
         buff = file_io.read(BLOCKSIZE)
 
         while len(buff) > 0:
             if isinstance(buff, str):
-                buff = buff.encode('utf-8')
+                buff = buff.encode("utf-8")
 
             hasher.update(buff)
             buff = file_io.read(BLOCKSIZE)
@@ -507,24 +432,22 @@ def get_file_sha1(filename_or_io):
     return digest
 
 
-def read_buffer(type_, io, output_queue, print_output=False, print_func=None):
-    '''
-    Reads a file-like buffer object into lines and optionally prints the output.
-    '''
-
-    def _print(line):
-        if print_func:
-            line = print_func(line)
+def get_path_permissions_mode(pathname: str):
+    """
+    Get the permissions (bits) of a path as an integer.
+    """
 
-        print(line)
+    mode_octal = oct(stat(pathname).st_mode)
+    return int(mode_octal[-3:])
 
-    for line in io:
-        # Handle local Popen shells returning list of bytes, not strings
-        if not isinstance(line, str):
-            line = line.decode('utf-8')
 
-        line = line.strip()
-        output_queue.put((type_, line))
-
-        if print_output:
-            _print(line)
+def raise_if_bad_type(
+    value: Any,
+    type_: Type,
+    exception: type[Exception],
+    message_prefix: str,
+):
+    try:
+        check_type(value, type_)
+    except TypeCheckError as e:
+        raise exception(f"{message_prefix}: {e}")