Provide revision post-write hooks

Added "post write hooks" to revision generation.   The primary rationale is
to provide for code formatting tools to automatically format new revisions,
however any arbitrary script or Python function may be  invoked as a hook.
The hooks are enabled by providing a ``[post_write_hooks]`` section in the
alembic.ini file.   The provided hook is a command-line runner which
includes configuration examples for running Black or autopep8 on newly
generated revision scripts.

The documentation also illustrates a custom hook that converts Python
source spaces to tabs, as requested in #577.

Co-authored-by: Mike Bayer <mike_mp@zzzcomputing.com>
Fixes: #307
Fixes: #577

Change-Id: I9d2092d20ec23f62ed3b33d979c16b979a450b48

2 files changed, 123 insertions, 1 deletions

diff --git a/alembic/script/base.py b/alembic/script/base.py
index b386dea..8aff3b5 100644
--- a/alembic/script/base.py
+++ b/alembic/script/base.py
@@ -7,6 +7,7 @@ import shutil
 from dateutil import tz
 
 from . import revision
+from . import write_hooks
 from .. import util
 from ..runtime import migration
 from ..util import compat
@@ -17,7 +18,7 @@ _legacy_rev = re.compile(r"([a-f0-9]+)\.py$")
 _mod_def_re = re.compile(r"(upgrade|downgrade)_([a-z0-9]+)")
 _slug_re = re.compile(r"\w+")
 _default_file_template = "%(rev)s_%(slug)s"
-_split_on_space_comma = re.compile(r",|(?: +)")
+_split_on_space_comma = re.compile(r", *|(?: +)")
 
 
 class ScriptDirectory(object):
@@ -50,6 +51,7 @@ class ScriptDirectory(object):
         sourceless=False,
         output_encoding="utf-8",
         timezone=None,
+        hook_config=None,
     ):
         self.dir = dir
         self.file_template = file_template
@@ -59,6 +61,7 @@ class ScriptDirectory(object):
         self.output_encoding = output_encoding
         self.revision_map = revision.RevisionMap(self._load_revisions)
         self.timezone = timezone
+        self.hook_config = hook_config
 
         if not os.access(dir, os.F_OK):
             raise util.CommandError(
@@ -143,6 +146,7 @@ class ScriptDirectory(object):
             output_encoding=config.get_main_option("output_encoding", "utf-8"),
             version_locations=version_locations,
             timezone=config.get_main_option("timezone"),
+            hook_config=config.get_section("post_write_hooks", {}),
         )
 
     @contextmanager
@@ -637,6 +641,11 @@ class ScriptDirectory(object):
             message=message if message is not None else ("empty message"),
             **kw
         )
+
+        post_write_hooks = self.hook_config
+        if post_write_hooks:
+            write_hooks._run_hooks(path, post_write_hooks)
+
         try:
             script = Script._from_path(self, path)
         except revision.RevisionError as err:
diff --git a/alembic/script/write_hooks.py b/alembic/script/write_hooks.py
new file mode 100644
index 0000000..61a6a27
--- /dev/null
+++ b/alembic/script/write_hooks.py
@@ -0,0 +1,113 @@
+import subprocess
+import sys
+
+from .. import util
+from ..util import compat
+
+
+_registry = {}
+
+
+def register(name):
+    """A function decorator that will register that function as a write hook.
+
+    See the documentation linked below for an example.
+
+    .. versionadded:: 1.2.0
+
+    .. seealso::
+
+        :ref:`post_write_hooks_custom`
+
+
+    """
+
+    def decorate(fn):
+        _registry[name] = fn
+
+    return decorate
+
+
+def _invoke(name, revision, options):
+    """Invokes the formatter registered for the given name.
+
+    :param name: The name of a formatter in the registry
+    :param revision: A :class:`.MigrationRevision` instance
+    :param options: A dict containing kwargs passed to the
+        specified formatter.
+    :raises: :class:`alembic.util.CommandError`
+    """
+    try:
+        hook = _registry[name]
+    except KeyError:
+        compat.raise_from_cause(
+            util.CommandError("No formatter with name '%s' registered" % name)
+        )
+    else:
+        return hook(revision, options)
+
+
+def _run_hooks(path, hook_config):
+    """Invoke hooks for a generated revision.
+
+    """
+
+    from .base import _split_on_space_comma
+
+    names = _split_on_space_comma.split(hook_config.get("hooks", ""))
+
+    for name in names:
+        if not name:
+            continue
+        opts = {
+            key[len(name) + 1 :]: hook_config[key]
+            for key in hook_config
+            if key.startswith(name + ".")
+        }
+        opts["_hook_name"] = name
+        try:
+            type_ = opts["type"]
+        except KeyError:
+            compat.raise_from_cause(
+                util.CommandError(
+                    "Key %s.type is required for post write hook %r"
+                    % (name, name)
+                )
+            )
+        else:
+            util.status(
+                'Running post write hook "%s"' % name,
+                _invoke,
+                type_,
+                path,
+                opts,
+                newline=True,
+            )
+
+
+@register("console_scripts")
+def console_scripts(path, options):
+    import pkg_resources
+
+    try:
+        entrypoint_name = options["entrypoint"]
+    except KeyError:
+        compat.raise_from_cause(
+            util.CommandError(
+                "Key %s.entrypoint is required for post write hook %r"
+                % (options["_hook_name"], options["_hook_name"])
+            )
+        )
+    iter_ = pkg_resources.iter_entry_points("console_scripts", entrypoint_name)
+    impl = next(iter_)
+    options = options.get("options", "")
+    subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import %s; %s()"
+            % (impl.module_name, ".".join((impl.module_name,) + impl.attrs)),
+            path,
+        ]
+        + options.split()
+    )