1import subprocess
2import sys
3
4from .. import util
5from ..util import compat
6
7
8_registry = {}
9
10
11def register(name):
12    """A function decorator that will register that function as a write hook.
13
14    See the documentation linked below for an example.
15
16    .. versionadded:: 1.2.0
17
18    .. seealso::
19
20        :ref:`post_write_hooks_custom`
21
22
23    """
24
25    def decorate(fn):
26        _registry[name] = fn
27
28    return decorate
29
30
31def _invoke(name, revision, options):
32    """Invokes the formatter registered for the given name.
33
34    :param name: The name of a formatter in the registry
35    :param revision: A :class:`.MigrationRevision` instance
36    :param options: A dict containing kwargs passed to the
37        specified formatter.
38    :raises: :class:`alembic.util.CommandError`
39    """
40    try:
41        hook = _registry[name]
42    except KeyError:
43        compat.raise_from_cause(
44            util.CommandError("No formatter with name '%s' registered" % name)
45        )
46    else:
47        return hook(revision, options)
48
49
50def _run_hooks(path, hook_config):
51    """Invoke hooks for a generated revision.
52
53    """
54
55    from .base import _split_on_space_comma
56
57    names = _split_on_space_comma.split(hook_config.get("hooks", ""))
58
59    for name in names:
60        if not name:
61            continue
62        opts = {
63            key[len(name) + 1 :]: hook_config[key]
64            for key in hook_config
65            if key.startswith(name + ".")
66        }
67        opts["_hook_name"] = name
68        try:
69            type_ = opts["type"]
70        except KeyError:
71            compat.raise_from_cause(
72                util.CommandError(
73                    "Key %s.type is required for post write hook %r"
74                    % (name, name)
75                )
76            )
77        else:
78            util.status(
79                'Running post write hook "%s"' % name,
80                _invoke,
81                type_,
82                path,
83                opts,
84                newline=True,
85            )
86
87
88@register("console_scripts")
89def console_scripts(path, options):
90    import pkg_resources
91
92    try:
93        entrypoint_name = options["entrypoint"]
94    except KeyError:
95        compat.raise_from_cause(
96            util.CommandError(
97                "Key %s.entrypoint is required for post write hook %r"
98                % (options["_hook_name"], options["_hook_name"])
99            )
100        )
101    iter_ = pkg_resources.iter_entry_points("console_scripts", entrypoint_name)
102    impl = next(iter_)
103    options = options.get("options", "")
104    subprocess.run(
105        [
106            sys.executable,
107            "-c",
108            "import %s; %s()"
109            % (impl.module_name, ".".join((impl.module_name,) + impl.attrs)),
110            path,
111        ]
112        + options.split()
113    )
114