codekingpro/portable-devtools

1# mypy: allow-untyped-defs, allow-incomplete-defs, allow-untyped-calls2# mypy: no-warn-return-any, allow-any-generics3 4from __future__ import annotations5 6import shlex7import subprocess8import sys9from typing import Any10from typing import Callable11from typing import Dict12from typing import List13from typing import Mapping14from typing import Optional15from typing import Union16 17from .. import util18from ..util import compat19 20 21REVISION_SCRIPT_TOKEN = "REVISION_SCRIPT_FILENAME"22 23_registry: dict = {}24 25 26def register(name: str) -> Callable:27    """A function decorator that will register that function as a write hook.28 29    See the documentation linked below for an example.30 31    .. seealso::32 33        :ref:`post_write_hooks_custom`34 35 36    """37 38    def decorate(fn):39        _registry[name] = fn40        return fn41 42    return decorate43 44 45def _invoke(46    name: str, revision: str, options: Mapping[str, Union[str, int]]47) -> Any:48    """Invokes the formatter registered for the given name.49 50    :param name: The name of a formatter in the registry51    :param revision: A :class:`.MigrationRevision` instance52    :param options: A dict containing kwargs passed to the53        specified formatter.54    :raises: :class:`alembic.util.CommandError`55    """56    try:57        hook = _registry[name]58    except KeyError as ke:59        raise util.CommandError(60            f"No formatter with name '{name}' registered"61        ) from ke62    else:63        return hook(revision, options)64 65 66def _run_hooks(path: str, hook_config: Mapping[str, str]) -> None:67    """Invoke hooks for a generated revision."""68 69    from .base import _split_on_space_comma70 71    names = _split_on_space_comma.split(hook_config.get("hooks", ""))72 73    for name in names:74        if not name:75            continue76        opts = {77            key[len(name) + 1 :]: hook_config[key]78            for key in hook_config79            if key.startswith(name + ".")80        }81        opts["_hook_name"] = name82        try:83            type_ = opts["type"]84        except KeyError as ke:85            raise util.CommandError(86                f"Key {name}.type is required for post write hook {name!r}"87            ) from ke88        else:89            with util.status(90                f"Running post write hook {name!r}", newline=True91            ):92                _invoke(type_, path, opts)93 94 95def _parse_cmdline_options(cmdline_options_str: str, path: str) -> List[str]:96    """Parse options from a string into a list.97 98    Also substitutes the revision script token with the actual filename of99    the revision script.100 101    If the revision script token doesn't occur in the options string, it is102    automatically prepended.103    """104    if REVISION_SCRIPT_TOKEN not in cmdline_options_str:105        cmdline_options_str = REVISION_SCRIPT_TOKEN + " " + cmdline_options_str106    cmdline_options_list = shlex.split(107        cmdline_options_str, posix=compat.is_posix108    )109    cmdline_options_list = [110        option.replace(REVISION_SCRIPT_TOKEN, path)111        for option in cmdline_options_list112    ]113    return cmdline_options_list114 115 116@register("console_scripts")117def console_scripts(118    path: str, options: dict, ignore_output: bool = False119) -> None:120    try:121        entrypoint_name = options["entrypoint"]122    except KeyError as ke:123        raise util.CommandError(124            f"Key {options['_hook_name']}.entrypoint is required for post "125            f"write hook {options['_hook_name']!r}"126        ) from ke127    for entry in compat.importlib_metadata_get("console_scripts"):128        if entry.name == entrypoint_name:129            impl: Any = entry130            break131    else:132        raise util.CommandError(133            f"Could not find entrypoint console_scripts.{entrypoint_name}"134        )135    cwd: Optional[str] = options.get("cwd", None)136    cmdline_options_str = options.get("options", "")137    cmdline_options_list = _parse_cmdline_options(cmdline_options_str, path)138 139    kw: Dict[str, Any] = {}140    if ignore_output:141        kw["stdout"] = kw["stderr"] = subprocess.DEVNULL142 143    subprocess.run(144        [145            sys.executable,146            "-c",147            f"import {impl.module}; {impl.module}.{impl.attr}()",148        ]149        + cmdline_options_list,150        cwd=cwd,151        **kw,152    )153 154 155@register("exec")156def exec_(path: str, options: dict, ignore_output: bool = False) -> None:157    try:158        executable = options["executable"]159    except KeyError as ke:160        raise util.CommandError(161            f"Key {options['_hook_name']}.executable is required for post "162            f"write hook {options['_hook_name']!r}"163        ) from ke164    cwd: Optional[str] = options.get("cwd", None)165    cmdline_options_str = options.get("options", "")166    cmdline_options_list = _parse_cmdline_options(cmdline_options_str, path)167 168    kw: Dict[str, Any] = {}169    if ignore_output:170        kw["stdout"] = kw["stderr"] = subprocess.DEVNULL171 172    subprocess.run(173        [174            executable,175            *cmdline_options_list,176        ],177        cwd=cwd,178        **kw,179    )180