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

12 files changed, 538 insertions, 6 deletions

diff --git a/alembic/config.py b/alembic/config.py
index 745ca8b..99c1ca8 100644
--- a/alembic/config.py
+++ b/alembic/config.py
@@ -212,11 +212,14 @@ class Config(object):
         """
         return os.path.join(package_dir, "templates")
 
-    def get_section(self, name):
+    def get_section(self, name, default=None):
         """Return all the configuration options from a given .ini file section
         as a dictionary.
 
         """
+        if not self.file_config.has_section(name):
+            return default
+
         return dict(self.file_config.items(name))
 
     def set_main_option(self, name, value):
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()
+    )
diff --git a/alembic/templates/generic/alembic.ini.mako b/alembic/templates/generic/alembic.ini.mako
index 1b7d6ea..281794f 100644
--- a/alembic/templates/generic/alembic.ini.mako
+++ b/alembic/templates/generic/alembic.ini.mako
@@ -38,6 +38,17 @@ script_location = ${script_location}
 sqlalchemy.url = driver://user:pass@localhost/dbname
 
 
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks=black
+# black.type=console_scripts
+# black.entrypoint=black
+# black.options=-l 79
+
 # Logging configuration
 [loggers]
 keys = root,sqlalchemy,alembic
diff --git a/alembic/templates/multidb/alembic.ini.mako b/alembic/templates/multidb/alembic.ini.mako
index 79fcb79..0b0919e 100644
--- a/alembic/templates/multidb/alembic.ini.mako
+++ b/alembic/templates/multidb/alembic.ini.mako
@@ -43,6 +43,16 @@ sqlalchemy.url = driver://user:pass@localhost/dbname
 [engine2]
 sqlalchemy.url = driver://user:pass@localhost/dbname2
 
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks=black
+# black.type=console_scripts
+# black.entrypoint=black
+# black.options=-l 79
 
 # Logging configuration
 [loggers]
diff --git a/alembic/templates/pylons/alembic.ini.mako b/alembic/templates/pylons/alembic.ini.mako
index 6f6511b..70fead0 100644
--- a/alembic/templates/pylons/alembic.ini.mako
+++ b/alembic/templates/pylons/alembic.ini.mako
@@ -35,6 +35,17 @@ script_location = ${script_location}
 # are written from script.py.mako
 # output_encoding = utf-8
 
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks=black
+# black.type=console_scripts
+# black.entrypoint=black
+# black.options=-l 79
+
 pylons_config_file = ./development.ini
 
 # that's it !
\ No newline at end of file
diff --git a/alembic/util/messaging.py b/alembic/util/messaging.py
index c91eb09..65b92c8 100644
--- a/alembic/util/messaging.py
+++ b/alembic/util/messaging.py
@@ -46,13 +46,14 @@ def write_outstream(stream, *text):
 
 
 def status(_statmsg, fn, *arg, **kw):
-    msg(_statmsg + " ...", False)
+    newline = kw.pop("newline", False)
+    msg(_statmsg + " ...", newline, True)
     try:
         ret = fn(*arg, **kw)
-        write_outstream(sys.stdout, " done\n")
+        write_outstream(sys.stdout, "  done\n")
         return ret
     except:
-        write_outstream(sys.stdout, " FAILED\n")
+        write_outstream(sys.stdout, "  FAILED\n")
         raise
 
 
@@ -73,7 +74,7 @@ def warn(msg, stacklevel=2):
     warnings.warn(msg, UserWarning, stacklevel=stacklevel)
 
 
-def msg(msg, newline=True):
+def msg(msg, newline=True, flush=False):
     if TERMWIDTH is None:
         write_outstream(sys.stdout, msg)
         if newline:
@@ -85,6 +86,8 @@ def msg(msg, newline=True):
             for line in lines[0:-1]:
                 write_outstream(sys.stdout, "  ", line, "\n")
         write_outstream(sys.stdout, "  ", lines[-1], ("\n" if newline else ""))
+    if flush:
+        sys.stdout.flush()
 
 
 def format_as_comma(value):
diff --git a/docs/build/api/script.rst b/docs/build/api/script.rst
index 8dc594b..cf58bb8 100644
--- a/docs/build/api/script.rst
+++ b/docs/build/api/script.rst
@@ -18,3 +18,9 @@ management, used exclusively by :class:`.ScriptDirectory`.
 
 .. automodule:: alembic.script.revision
     :members:
+
+Write Hooks
+===========
+
+.. automodule:: alembic.script.write_hooks
+    :members:
diff --git a/docs/build/autogenerate.rst b/docs/build/autogenerate.rst
index 3a9e0c9..66d054e 100644
--- a/docs/build/autogenerate.rst
+++ b/docs/build/autogenerate.rst
@@ -489,4 +489,177 @@ then a basic check for type equivalence is run.
    method.
 
 
+.. _post_write_hooks:
 
+Applying Post Processing and Python Code Formatters to Generated Revisions
+---------------------------------------------------------------------------
+
+Revision scripts generated by the ``alembic revision`` command can optionally
+be piped through a series of post-production functions which may analyze or
+rewrite Python source code generated by Alembic, within the scope of running
+the ``revision`` command.   The primary intended use of this feature is to run
+code-formatting tools such as `Black <https://black.readthedocs.io/>`_ or
+`autopep8 <https://pypi.org/project/autopep8/>`_, as well as custom-written
+formatting and linter functions, on revision files as Alembic generates them.
+Any number of hooks can be configured and they will be run in series, given the
+path to the newly generated file as well as configuration options.
+
+The post write hooks, when configured,  run against generated revision files
+regardless of whether or not the autogenerate feature was used.
+
+.. versionadded:: 1.2
+
+.. note::
+
+    Alembic's post write system is partially inspired by the `pre-commit
+    <https://pre-commit.com/>`_ tool, which configures git hooks that reformat
+    source files as they are committed to a git repository.  Pre-commit can
+    serve this role for Alembic revision files as well, applying code
+    formatters to them as they are committed.  Alembic's post write hooks are
+    useful only in that they can format the files immediately upon generation,
+    rather than at commit time, and also can be useful for projects that prefer
+    not to use post-commit.
+
+
+Basic Formatter Configuration
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The ``alembic.ini`` samples now include commented-out configuration
+illustrating how to configure code-formatting tools to run against the newly
+generated file path.    Example::
+
+  [post_write_hooks]
+
+  # format using "black"
+  hooks=black
+
+  black.type=console_scripts
+  black.entrypoint=black
+  black.options=-l 79
+
+Above, we configure a single post write hook that we call ``"black"``. Note
+that this name is arbitrary.  We then define the configuration for the
+``"black"`` post write hook, which includes:
+
+* ``type`` - this is the type of hook we are running.   Alembic includes
+  a hook runner called ``"console_scripts"``, which is specifically a
+  Python function that uses ``subprocess.run()`` to invoke a separate
+  Python script against the revision file.    For a custom-written hook
+  function, this configuration variable would refer to the name under
+  which the custom hook was registered; see the next section for an example.
+
+* ``entrypoint`` - this part of the configuration is specific to the
+  ``"console_scripts"`` hook runner.  This is the name of the `setuptools entrypoint <https://setuptools.readthedocs.io/en/latest/pkg_resources.html#entry-points>`_
+  that is used to define the console script.    Within the scope of standard
+  Python console scripts, this name will match the name of the shell command
+  that is usually run for the code formatting tool, in this case ``black``.
+
+* ``options`` - this is also specific to the ``"console_scripts"`` hook runner.
+  This is a line of command-line options that will be passed to the
+  code formatting tool.  In this case, we want to run the command
+  as ``black -l 79 /path/to/revision.py``.   The path of the revision file
+  is sent as a single positional argument to the script after the options.
+
+  .. note:: Make sure options for the script are provided such that it will
+     rewrite the input file **in place**.  For example, when running
+     ``autopep8``, the ``--in-place`` option should be provided::
+
+        [post_write_hooks]
+        hooks=autopep8
+        autopep8.type=console_scripts
+        autopep8.entrypoint=autopep8
+        autopep8.options=--in-place
+
+
+When running ``alembic revision -m "rev1"``, we will now see the ``black``
+tool's output as well::
+
+  $ alembic revision -m "rev1"
+    Generating /path/to/project/versions/481b13bc369a_rev1.py ... done
+    Running post write hook "black" ...
+  reformatted /path/to/project/versions/481b13bc369a_rev1.py
+  All done! ✨ 🍰 ✨
+  1 file reformatted.
+    done
+
+Hooks may also be specified as a list of names, which correspond to hook
+runners that will run sequentially.  As an example, we can also run the
+`zimports <https://pypi.org/project/zimports/>`_ import rewriting tool (written
+by Alembic's author) subsequent to running the ``black`` tool, using a
+configuration as follows::
+
+  [post_write_hooks]
+
+  # format using "black", then "zimports"
+  hooks=black, zimports
+
+  black.type=console_scripts
+  black.entrypoint=black
+  black.options=-l 79
+
+  zimports.type=console_scripts
+  zimports.entrypoint=zimports
+  zimports.options=--style google
+
+When using the above configuration, a newly generated revision file will
+be processed first by the "black" tool, then by the "zimports" tool.
+
+.. _post_write_hooks_custom:
+
+Writing Custom Hooks as Python Functions
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The previous section illustrated how to run command-line code formatters,
+through the use of a post write hook provided by Alembic known as
+``console_scripts``.  This hook is in fact a Python function that is registered
+under that name using a registration function that may be used to register
+other types of hooks as well.
+
+To illustrate, we will use the example of a short Python function that wants
+to rewrite the generated code to use tabs instead of four spaces.   For simplicity,
+we will illustrate how this function can be present directly in the ``env.py``
+file.   The function is declared and registered using the :func:`.write_hooks.register`
+decorator::
+
+    from alembic.script import write_hooks
+    import re
+
+    @write_hooks.register("spaces_to_tabs")
+    def convert_spaces_to_tabs(filename, options):
+        lines = []
+        with open(filename) as file_:
+            for line in file_:
+                lines.append(
+                    re.sub(
+                        r"^(    )+",
+                        lambda m: "\t" * (len(m.group(1)) // 4),
+                        line
+                    )
+                )
+        with open(filename, "w") as to_write:
+            to_write.write("".join(lines))
+
+Our new ``"spaces_to_tabs"`` hook can be configured in alembic.ini as follows::
+
+  [alembic]
+
+  # ...
+
+  # ensure the revision command loads env.py
+  revision_environment = true
+
+  [post_write_hooks]
+
+  hooks=spaces_to_tabs
+
+  spaces_to_tabs.type=spaces_to_tabs
+
+
+When ``alembic revision`` is run, the ``env.py`` file will be loaded in all
+cases, the custom "spaces_to_tabs" function will be registered and it will then
+be run against the newly generated file path::
+
+  $ alembic revision -m "rev1"
+    Generating /path/to/project/versions/481b13bc369a_rev1.py ... done
+    Running post write hook "spaces_to_tabs" ...
+    done
diff --git a/docs/build/tutorial.rst b/docs/build/tutorial.rst
index 458717f..f3d513c 100644
--- a/docs/build/tutorial.rst
+++ b/docs/build/tutorial.rst
@@ -149,6 +149,17 @@ The file generated with the "generic" configuration looks like::
 
     sqlalchemy.url = driver://user:pass@localhost/dbname
 
+    # post_write_hooks defines scripts or Python functions that are run
+    # on newly generated revision scripts.  See the documentation for further
+    # detail and examples
+
+    # format using "black" - use the console_scripts runner,
+    # against the "black" entrypoint
+    # hooks=black
+    # black.type=console_scripts
+    # black.entrypoint=black
+    # black.options=-l 79
+
     # Logging configuration
     [loggers]
     keys = root,sqlalchemy,alembic
diff --git a/docs/build/unreleased/307.rst b/docs/build/unreleased/307.rst
new file mode 100644
index 0000000..988e7a4
--- /dev/null
+++ b/docs/build/unreleased/307.rst
@@ -0,0 +1,18 @@
+ .. change::
+    :tags: feature, commands
+    :tickets: 307
+
+    Added "post write hooks" to revision generation.  These allow custom logic
+    to run after a revision Python script is generated, typically for the
+    purpose of running code formatters such as "Black" or "autopep8", but may
+    be used for any arbitrary post-render hook as well, including custom Python
+    functions or scripts.  The hooks are enabled by providing a
+    ``[post_write_hooks]`` section in the alembic.ini file.  A single hook
+    is provided which runs an arbitrary Python executable on the newly
+    generated revision script, which can be configured to run code formatters
+    such as Black; full examples are included in the documentation.
+
+    .. seealso::
+
+        :ref:`post_write_hooks`
+
diff --git a/tests/test_post_write.py b/tests/test_post_write.py
new file mode 100644
index 0000000..f93ca90
--- /dev/null
+++ b/tests/test_post_write.py
@@ -0,0 +1,164 @@
+import sys
+
+from alembic import command
+from alembic import util
+from alembic.script import write_hooks
+from alembic.testing import assert_raises_message
+from alembic.testing import eq_
+from alembic.testing import mock
+from alembic.testing import TestBase
+from alembic.testing.env import _no_sql_testing_config
+from alembic.testing.env import clear_staging_env
+from alembic.testing.env import staging_env
+
+
+class HookTest(TestBase):
+    def test_register(self):
+        @write_hooks.register("my_writer")
+        def my_writer(path, config):
+            return path
+
+        assert "my_writer" in write_hooks._registry
+
+    def test_invoke(self):
+        my_formatter = mock.Mock()
+        write_hooks.register("my_writer")(my_formatter)
+
+        write_hooks._invoke("my_writer", "/some/path", {"option": 1})
+
+        my_formatter.assert_called_once_with("/some/path", {"option": 1})
+
+
+class RunHookTest(TestBase):
+    def setUp(self):
+        self.env = staging_env()
+
+    def tearDown(self):
+        clear_staging_env()
+
+    def test_generic(self):
+        hook1 = mock.Mock()
+        hook2 = mock.Mock()
+
+        write_hooks.register("hook1")(hook1)
+        write_hooks.register("hook2")(hook2)
+
+        self.cfg = _no_sql_testing_config(
+            directives=(
+                "\n[post_write_hooks]\n"
+                "hooks=hook1,hook2\n"
+                "hook1.type=hook1\n"
+                "hook1.arg1=foo\n"
+                "hook2.type=hook2\n"
+                "hook2.arg1=bar\n"
+            )
+        )
+
+        rev = command.revision(self.cfg, message="x")
+
+        eq_(
+            hook1.mock_calls,
+            [
+                mock.call(
+                    rev.path,
+                    {"type": "hook1", "arg1": "foo", "_hook_name": "hook1"},
+                )
+            ],
+        )
+        eq_(
+            hook2.mock_calls,
+            [
+                mock.call(
+                    rev.path,
+                    {"type": "hook2", "arg1": "bar", "_hook_name": "hook2"},
+                )
+            ],
+        )
+
+    def test_empty_section(self):
+        self.cfg = _no_sql_testing_config(
+            directives=("\n[post_write_hooks]\n")
+        )
+
+        command.revision(self.cfg, message="x")
+
+    def test_no_section(self):
+        self.cfg = _no_sql_testing_config(directives="")
+
+        command.revision(self.cfg, message="x")
+
+    def test_empty_hooks(self):
+        self.cfg = _no_sql_testing_config(
+            directives=("\n[post_write_hooks]\n" "hooks=\n")
+        )
+
+        command.revision(self.cfg, message="x")
+
+    def test_no_type(self):
+        self.cfg = _no_sql_testing_config(
+            directives=(
+                "\n[post_write_hooks]\n" "hooks=foo\n" "foo.bar=somebar\n"
+            )
+        )
+
+        assert_raises_message(
+            util.CommandError,
+            "Key foo.type is required for post write hook 'foo'",
+            command.revision,
+            self.cfg,
+            message="x",
+        )
+
+    def test_console_scripts_entrypoint_missing(self):
+        self.cfg = _no_sql_testing_config(
+            directives=(
+                "\n[post_write_hooks]\n"
+                "hooks=black\n"
+                "black.type=console_scripts\n"
+            )
+        )
+        assert_raises_message(
+            util.CommandError,
+            "Key black.entrypoint is required for post write hook 'black'",
+            command.revision,
+            self.cfg,
+            message="x",
+        )
+
+    def test_console_scripts(self):
+        self.cfg = _no_sql_testing_config(
+            directives=(
+                "\n[post_write_hooks]\n"
+                "hooks=black\n"
+                "black.type=console_scripts\n"
+                "black.entrypoint=black\n"
+                "black.options=-l 79\n"
+            )
+        )
+
+        impl = mock.Mock(attrs=("foo", "bar"), module_name="black_module")
+        entrypoints = mock.Mock(return_value=iter([impl]))
+        with mock.patch(
+            "pkg_resources.iter_entry_points", entrypoints
+        ), mock.patch(
+            "alembic.script.write_hooks.subprocess"
+        ) as mock_subprocess:
+
+            rev = command.revision(self.cfg, message="x")
+
+        eq_(entrypoints.mock_calls, [mock.call("console_scripts", "black")])
+        eq_(
+            mock_subprocess.mock_calls,
+            [
+                mock.call.run(
+                    [
+                        sys.executable,
+                        "-c",
+                        "import black_module; black_module.foo.bar()",
+                        rev.path,
+                        "-l",
+                        "79",
+                    ]
+                )
+            ],
+        )