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

pyinfra_cli/util.py

@@ -1,98 +1,143 @@
-import json
+from __future__ import annotations
 
+import json
+import os
 from datetime import datetime
 from importlib import import_module
+from importlib.util import find_spec
 from io import IOBase
-from types import FunctionType
+from os import path
+from pathlib import Path
+from types import CodeType, FunctionType, ModuleType
+from typing import Callable
 
 import click
+import gevent
+
+from pyinfra import logger, state
+from pyinfra.api.command import PyinfraCommand
+from pyinfra.api.exceptions import PyinfraError
+from pyinfra.api.host import HostData
+from pyinfra.api.operation import OperationMeta
+from pyinfra.api.state import (
+    State,
+    StateHostMeta,
+    StateHostResults,
+    StateOperationHostData,
+    StateOperationMeta,
+)
+from pyinfra.context import ctx_config, ctx_host
+from pyinfra.progress import progress_spinner
+
+from .exceptions import CliError, UnexpectedExternalError
 
-from pyinfra import logger, pseudo_host, pseudo_state
-from pyinfra.api.util import FallbackDict
-from pyinfra.hook import HOOKS
+# Cache for compiled Python deploy code
+PYTHON_CODES: dict[str, CodeType] = {}
 
-from .exceptions import CliError
 
-# Cache for compiled Python deploy code
-PYTHON_CODES = {}
+def is_subdir(child, parent):
+    child = path.realpath(child)
+    parent = path.realpath(parent)
+    relative = path.relpath(child, start=parent)
+    return not relative.startswith(os.pardir)
 
 
-def exec_file(filename, return_locals=False, is_deploy_code=False):
-    '''
+def exec_file(filename, return_locals: bool = False, is_deploy_code: bool = False):
+    """
     Execute a Python file and optionally return it's attributes as a dict.
-    '''
+    """
+
+    old_current_exec_filename = state.current_exec_filename
+    state.current_exec_filename = filename
 
     if filename not in PYTHON_CODES:
-        with open(filename, 'r') as f:
-            code = f.read()
+        with open(filename, "r", encoding="utf-8") as f:
+            code_str = f.read()
 
-        code = compile(code, filename, 'exec')
+        code = compile(code_str, filename, "exec")
         PYTHON_CODES[filename] = code
 
     # Create some base attributes for our "module"
-    data = {
-        '__file__': filename,
-        'state': pseudo_state,
-    }
+    data = {"__file__": filename}
 
     # Execute the code with locals/globals going into the dict above
-    exec(PYTHON_CODES[filename], data)
+    try:
+        exec(PYTHON_CODES[filename], data)
+    except PyinfraError:
+        # Raise pyinfra errors as-is
+        raise
+    except Exception as e:
+        # Wrap & re-raise errors in user code so we highlight filename/etc
+        raise UnexpectedExternalError(e, filename)
+    finally:
+        state.current_exec_filename = old_current_exec_filename
 
     return data
 
 
-def run_hook(state, hook_name, hook_data):
-    hooks = HOOKS[hook_name]
-
-    if hooks:
-        for hook in hooks:
-            print('--> Running hook: {0}/{1}'.format(
-                hook_name,
-                click.style(hook.__name__, bold=True),
-            ))
-            hook(hook_data, state)
-
-        print()
+def json_encode(obj):
+    # pyinfra types
+    if isinstance(obj, HostData):
+        return obj.dict()
 
+    if isinstance(obj, PyinfraCommand):
+        return repr(obj)
+
+    if isinstance(
+        obj,
+        (
+            OperationMeta,
+            StateOperationMeta,
+            StateOperationHostData,
+            StateHostMeta,
+            StateHostResults,
+        ),
+    ):
+        return repr(obj)
+
+    # Python types
+    if isinstance(obj, ModuleType):
+        return "Module: {0}".format(obj.__name__)
 
-def json_encode(obj):
     if isinstance(obj, FunctionType):
-        return obj.__name__
+        return "Function: {0}".format(obj.__name__)
 
-    elif isinstance(obj, datetime):
+    if isinstance(obj, datetime):
         return obj.isoformat()
 
-    elif isinstance(obj, FallbackDict):
-        return obj.dict()
-
-    elif isinstance(obj, IOBase):
-        if hasattr(obj, 'name'):
-            return 'File: {0}'.format(obj.name)
+    if isinstance(obj, IOBase):
+        if hasattr(obj, "name"):
+            return "File: {0}".format(obj.name)
 
-        elif hasattr(obj, 'template'):
-            return 'Template: {0}'.format(obj.template)
+        if hasattr(obj, "template"):
+            return "Template: {0}".format(obj.template)
 
         obj.seek(0)
-        return 'In memory file: {0}'.format(obj.read())
+        return "In memory file: {0}".format(obj.read())
+
+    if isinstance(obj, Path):
+        return str(obj)
 
-    elif isinstance(obj, set):
+    if isinstance(obj, set):
         return sorted(list(obj))
 
-    elif isinstance(obj, bytes):
+    if isinstance(obj, bytes):
         return obj.decode()
 
-    else:
-        raise TypeError('Cannot serialize: {0} ({1})'.format(type(obj), obj))
+    if hasattr(obj, "to_json"):
+        return obj.to_json()
+
+    raise TypeError("Cannot serialize: {0} ({1})".format(type(obj), obj))
 
 
-def _parse_arg(arg):
+def parse_cli_arg(arg):
     if isinstance(arg, list):
-        return [_parse_arg(a) for a in arg]
+        return [parse_cli_arg(a) for a in arg]
 
-    if arg.lower() == 'false':
+    if arg.lower() == "false":
         return False
 
-    if arg.lower() == 'true':
+    if arg.lower() == "true":
         return True
 
     try:
@@ -100,73 +145,95 @@ def _parse_arg(arg):
     except (TypeError, ValueError):
         pass
 
+    try:
+        return json.loads(arg)
+    except ValueError:
+        pass
+
     return arg
 
 
-def get_operation_and_args(commands):
-    operation_name = commands[0]
+def try_import_module_attribute(path, prefix=None, raise_for_none=True):
+    if ":" in path:
+        # Allow a.module.name:function syntax
+        mod_path, attr_name = path.rsplit(":", 1)
+    elif "." in path:
+        # And also a.module.name.function
+        mod_path, attr_name = path.rsplit(".", 1)
+    else:
+        return None
 
-    # Get the module & operation name
-    op_module, op_name = operation_name.split('.')
+    possible_modules = [mod_path]
+    if prefix:
+        possible_modules.append(f"{prefix}.{mod_path}")
 
-    try:
-        op_module = import_module('pyinfra.modules.{0}'.format(op_module))
-    except ImportError:
-        raise CliError('No such module: {0}'.format(op_module))
+    module = None
 
-    op = getattr(op_module, op_name, None)
-    if not op:
-        raise CliError('No such operation: {0}'.format(operation_name))
+    for possible in possible_modules:
+        try:
+            # Look for the module/fn, note that from the find_spec doc:
+            # "If the name is for submodule (contains a dot), the parent module is
+            # automatically imported."
+            spec = find_spec(possible)
+        except ModuleNotFoundError:
+            continue
+        except Exception as e:
+            # Capture all exceptions here which may be triggered from the automatic module import
+            # referenced above the find_spec call.
+            logger.warning(f"Exception raised during inventory search on: {possible}: {e}")
+            continue
+        else:
+            if spec is not None:
+                module = import_module(possible)
+                break
 
-    # Parse the arguments
-    operation_args = commands[1:]
+    if module is None:
+        if raise_for_none:
+            raise CliError(f"No such module: {possible_modules[0]}")
+        return
 
-    if len(operation_args) == 1:
-        # Check if we're JSON (in which case we expect a list of two items:
-        # a list of args and a dict of kwargs).
-        try:
-            args, kwargs = json.loads(operation_args[0])
-            return op, (args, kwargs)
-        except ValueError:
-            pass
-
-    args = [
-        _parse_arg(arg)
-        for arg in operation_args if '=' not in arg
-    ]
-
-    kwargs = {
-        key: _parse_arg(value)
-        for key, value in [
-            arg.split('=', 1)
-            for arg in operation_args if '=' in arg
-        ]
-    }
+    attr = getattr(module, attr_name, None)
+    if attr is None:
+        if raise_for_none:
+            raise CliError(f"No such attribute in module {possible_modules[0]}: {attr_name}")
+        return
 
-    return op, (args, kwargs)
+    return attr
 
 
-def load_deploy_file(state, filename):
-    # Copy the inventory hosts (some might be removed during deploy)
-    hosts = list(state.inventory)
+def _parallel_load_hosts(state: "State", callback: Callable, name: str):
+    def load_file(local_host):
+        try:
+            with ctx_config.use(state.config.copy()):
+                with ctx_host.use(local_host):
+                    callback()
+                    logger.info(
+                        "{0}{1} {2}".format(
+                            local_host.print_prefix,
+                            click.style("Ready:", "green"),
+                            click.style(name, bold=True),
+                        ),
+                    )
+        except Exception as e:
+            return e
+
+    greenlet_to_host = {
+        state.pool.spawn(load_file, host): host for host in state.inventory.iter_active_hosts()
+    }
 
-    for host in hosts:
-        # Don't load for anything within our (top level, --limit) limit
-        if (
-            isinstance(state.limit_hosts, list)
-            and host not in state.limit_hosts
-        ):
-            continue
+    with progress_spinner(greenlet_to_host.values()) as progress:
+        for greenlet in gevent.iwait(greenlet_to_host.keys()):
+            host = greenlet_to_host[greenlet]
+            result = greenlet.get()
+            if isinstance(result, Exception):
+                raise result
+            progress(host)
 
-        pseudo_host.set(host)
 
-        exec_file(filename)
+def load_deploy_file(state: "State", filename):
+    state.current_deploy_filename = filename
+    _parallel_load_hosts(state, lambda: exec_file(filename), filename)
 
-        logger.info('{0} {1} {2}'.format(
-            host.print_prefix,
-            click.style('Ready:', 'green'),
-            click.style(filename, bold=True),
-        ))
 
-    # Remove any pseudo host
-    pseudo_host.reset()
+def load_func(state: "State", op_func, *args, **kwargs):
+    _parallel_load_hosts(state, lambda: op_func(*args, **kwargs), op_func.__name__)

pyinfra_cli/virtualenv.py

@@ -1,11 +1,13 @@
 import os
 import sys
 
+import click
+
 from pyinfra import logger
 
 
-def init_virtualenv():
-    '''
+def init_virtualenv() -> None:
+    """
     Add a virtualenv to sys.path so the user can import modules from it.
     This isn't perfect: it doesn't use the Python interpreter with which the
     virtualenv was built, and it ignores the --no-site-packages option. A
@@ -14,14 +16,14 @@ def init_virtualenv():
 
     Adapted from IPython's implementation:
     https://github.com/ipython/ipython/blob/master/IPython/core/interactiveshell.py
-    '''
+    """
 
-    if 'VIRTUAL_ENV' not in os.environ:
+    if "VIRTUAL_ENV" not in os.environ:
         # Not in a virtualenv
         return
 
     p = os.path.normcase(sys.executable)
-    p_venv = os.path.normcase(os.environ['VIRTUAL_ENV'])
+    p_venv = os.path.normcase(os.environ["VIRTUAL_ENV"])
 
     # executable path should end like /bin/python or \\scripts\\python.exe
     p_exe_up2 = os.path.dirname(os.path.dirname(p))
@@ -39,35 +41,38 @@ def init_virtualenv():
         paths.append(p)
 
     # In Cygwin paths like 'c:\...' and '\cygdrive\c\...' are possible
-    if p_venv.startswith('\\cygdrive'):
+    if p_venv.startswith("\\cygdrive"):
         p_venv = p_venv[11:]
-    elif len(p_venv) >= 2 and p_venv[1] == ':':
+    elif len(p_venv) >= 2 and p_venv[1] == ":":
         p_venv = p_venv[2:]
 
     if any(p_venv in p for p in paths):
         # Running properly in the virtualenv, don't need to do anything
         return
 
-    logger.warning((
-        'Attempting to work in a virtualenv.\n'
-        '    If you encounter problems, please install pyinfra inside the virtualenv.'
-    ))
-    print()
+    logger.warning(
+        (
+            "Attempting to work in a virtualenv.\n"
+            "    If you encounter problems, please install pyinfra inside the virtualenv."
+        ),
+    )
+    click.echo(err=True)
 
-    if sys.platform == 'win32':
+    if sys.platform == "win32":
         virtual_env = os.path.join(
-            os.environ['VIRTUAL_ENV'],
-            'Lib',
-            'site-packages',
+            os.environ["VIRTUAL_ENV"],
+            "Lib",
+            "site-packages",
         )
     else:
         virtual_env = os.path.join(
-            os.environ['VIRTUAL_ENV'],
-            'lib',
-            'python%d.%d' % sys.version_info[:2],
-            'site-packages',
+            os.environ["VIRTUAL_ENV"],
+            "lib",
+            "python%d.%d" % sys.version_info[:2],
+            "site-packages",
         )
 
     import site
+
     sys.path.insert(0, virtual_env)
     site.addsitedir(virtual_env)

pyinfra/api/connectors/__init__.py

@@ -1,21 +0,0 @@
-from . import ansible, docker, local, ssh, vagrant
-
-
-# Connectors that handle execution of pyinfra operations
-EXECUTION_CONNECTORS = {  # pragma: no cover
-    'docker': docker,
-    'local': local,
-    'ssh': ssh,
-}
-
-# Connectors that handle generation of inventories
-INVENTORY_CONNECTORS = {  # pragma: no cover
-    'docker': docker,
-    'vagrant': vagrant,
-    'ansible': ansible,
-}
-
-ALL_CONNECTORS = (  # pragma: no cover
-    list(EXECUTION_CONNECTORS.keys())
-    + list(INVENTORY_CONNECTORS.keys())
-)

pyinfra/api/connectors/ansible.py

@@ -1,99 +0,0 @@
-import re
-
-from collections import defaultdict
-from configparser import ConfigParser
-from os import path
-
-from pyinfra import logger
-from pyinfra.api.exceptions import InventoryError
-from pyinfra.api.util import memoize
-
-
-@memoize
-def show_warning():
-    logger.warning('The @ansible connector is in Alpha!')
-
-
-@memoize
-def get_ansible_inventory(inventory_filename=None):
-    if not inventory_filename:  # pragma: no cover
-        raise InventoryError('No Ansible inventory filename provided!')
-
-    if not path.exists(inventory_filename):
-        raise InventoryError((
-            'Could not find Ansible inventory file: {0}'
-        ).format(inventory_filename))
-
-    logger.info('Parsing Ansible inventory...')
-
-    config = ConfigParser(
-        delimiters=(' '),  # we only handle the hostnames for now
-        allow_no_value=True,  # we don't by default have = values
-        interpolation=None,  # remove any Python interpolation
-    )
-    config.read(inventory_filename)
-    return config
-
-
-def _parse_ansible_hosts(hosts):
-    for host in hosts:
-        expand_match = re.search(r'\[[0-9:]+\]', host)
-        if expand_match:
-            expand_string = host[expand_match.start():expand_match.end()]
-            bits = expand_string[1:-1].split(':')  # remove the [] either side
-
-            zfill = 0
-            if bits[0].startswith('0'):
-                zfill = len(bits[0])
-
-            start, end = int(bits[0]), int(bits[1])
-            step = int(bits[2]) if len(bits) > 2 else 1
-
-            for n in range(start, end + 1, step):
-                number_as_string = '{0}'.format(n)
-                if zfill:
-                    number_as_string = number_as_string.zfill(zfill)
-
-                hostname = host.replace(expand_string, number_as_string)
-                yield hostname
-        else:
-            yield host
-
-
-def make_names_data(inventory_filename=None):
-    show_warning()
-
-    config = get_ansible_inventory(inventory_filename)
-
-    host_to_groups = defaultdict(set)
-    group_to_hosts = defaultdict(set)
-    hosts = []
-
-    # First pass - load hosts/groups of hosts
-    for section in config.sections():
-        if ':' in section:  # ignore :children and :vars sections this time
-            continue
-
-        options = config.options(section)
-        for host in _parse_ansible_hosts(options):
-            hosts.append(host)
-            host_to_groups[host].add(section)
-            group_to_hosts[section].add(host)
-
-    # Second pass - load any children groups
-    for section in config.sections():
-        if not section.endswith(':children'):  # we only support :children for now
-            continue
-
-        group_name = section.replace(':children', '')
-
-        options = config.options(section)
-        for sub_group_name in options:
-            sub_group_hosts = group_to_hosts[sub_group_name]
-            for host in sub_group_hosts:
-                host_to_groups[host].add(group_name)
-
-    return [
-        (host, {}, sorted(list(host_to_groups.get(host))))
-        for host in hosts
-    ]