Add "module" post-write hook

This hook type is almost identical to the console_scripts hook, except
it's running `python -m black` instead of using black's console_script.

It is mainly useful for tools without console scripts (e.g. ruff), but
has semantics closer to the console_scripts hook in that it finds the
ruff module available to the running interpreter instead of finding
an executable by path.

Fixes #1686

diff --git a/alembic/script/write_hooks.py b/alembic/script/write_hooks.py
index 50aefbe364502c24eed4cda425232a1454d61319..cf96e91d8af7b499dc9deaa2789923012e5c2d27 100644 (file)
--- a/alembic/script/write_hooks.py
+++ b/alembic/script/write_hooks.py
@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import importlib.util
 import os
 import shlex
 import subprocess
@@ -176,3 +177,36 @@ def exec_(path: str, options: dict, ignore_output: bool = False) -> None:
         cwd=cwd,
         **kw,
     )
+
+
+@register("module")
+def module(path: str, options: dict, ignore_output: bool = False) -> None:
+    try:
+        module_name = options["module"]
+    except KeyError as ke:
+        raise util.CommandError(
+            f"Key {options['_hook_name']}.module is required for post "
+            f"write hook {options['_hook_name']!r}"
+        ) from ke
+
+    if importlib.util.find_spec(module_name) is None:
+        raise util.CommandError(f"Could not find module {module_name}")
+
+    cwd: Optional[str] = options.get("cwd", None)
+    cmdline_options_str = options.get("options", "")
+    cmdline_options_list = _parse_cmdline_options(cmdline_options_str, path)
+
+    kw: Dict[str, Any] = {}
+    if ignore_output:
+        kw["stdout"] = kw["stderr"] = subprocess.DEVNULL
+
+    subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            module_name,
+        ]
+        + cmdline_options_list,
+        cwd=cwd,
+        **kw,
+    )

diff --git a/docs/build/autogenerate.rst b/docs/build/autogenerate.rst
index aca8683b719fd75c7ebf5a49543a82ae8681c700..2c7892c65cb563e1d24b82a49076b2965027e85e 100644 (file)
--- a/docs/build/autogenerate.rst
+++ b/docs/build/autogenerate.rst
@@ -721,12 +721,16 @@ Above, we configure ``hooks`` to be a single post write hook labeled
 configuration for the ``"black"`` post write hook, which includes:
 
 * ``type`` - this is the type of hook we are running.  Alembic includes
-  two hook runners: ``"console_scripts"``, which is specifically a
-  Python function that uses ``subprocess.run()`` to invoke a separate
-  Python script against the revision file; and ``"exec"``, which uses
-  ``subprocess.run()`` to execute an arbitrary binary.  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.
+  three hook runners:
+
+  * ``"console_scripts"``, which is specifically a Python function that uses
+    ``subprocess.run()`` to invoke a separate Python script against the revision file;
+  * ``"exec"``, which uses ``subprocess.run()`` to execute an arbitrary binary; and
+  * ``"module"``, which uses ``subprocess.run()`` to invoke a Python module directly.
+
+  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.
 
 .. versionadded:: 1.12 added new ``exec`` runner
 

diff --git a/tests/test_post_write.py b/tests/test_post_write.py
index b3c277659b46dd76279fb43a30bede5f3d7feb03..fe2afe97f5cb978dedddcaa1e1d5463ec9a79467 100644 (file)
--- a/tests/test_post_write.py
+++ b/tests/test_post_write.py
@@ -1,3 +1,4 @@
+import importlib.util
 from pathlib import Path
 import sys
 
@@ -477,3 +478,193 @@ ruff.cwd = /path/to/cwd
         self._run_ruff_with_config(
             input_config, expected_additional_arguments_fn, cwd="/path/to/cwd"
         )
+
+    def test_module_config_missing(self):
+        self.cfg = _no_sql_testing_config(
+            directives=(
+                """
+[post_write_hooks]
+hooks = ruff
+ruff.type = module
+                """
+            )
+        )
+        assert_raises_message(
+            util.CommandError,
+            "Key ruff.module is required for post write hook 'ruff'",
+            command.revision,
+            self.cfg,
+            message="x",
+        )
+
+    def test_module_not_found(self):
+        self.cfg = _no_sql_testing_config(
+            directives=(
+                """
+[post_write_hooks]
+hooks = ruff
+ruff.type = module
+ruff.module = ruff_not_found
+                """
+            )
+        )
+        assert_raises_message(
+            util.CommandError,
+            "Could not find module ruff_not_found",
+            command.revision,
+            self.cfg,
+            message="x",
+        )
+
+    def _run_black_module_with_config(
+        self,
+        input_config,
+        expected_additional_arguments_fn,
+        cwd=None,
+        use_toml=False,
+    ):
+        if use_toml:
+            self.cfg = _no_sql_pyproject_config(directives=input_config)
+        else:
+            self.cfg = _no_sql_testing_config(directives=input_config)
+
+        black_module_spec = importlib.util.find_spec("black")
+
+        importlib_util_find_spec = mock.Mock(return_value=black_module_spec)
+        with (
+            mock.patch(
+                "importlib.util.find_spec",
+                importlib_util_find_spec,
+            ),
+            mock.patch(
+                "alembic.script.write_hooks.subprocess"
+            ) as mock_subprocess,
+        ):
+            rev = command.revision(self.cfg, message="x")
+
+        eq_(importlib_util_find_spec.mock_calls, [mock.call("black")])
+        eq_(
+            mock_subprocess.mock_calls,
+            [
+                mock.call.run(
+                    [
+                        sys.executable,
+                        "-m",
+                        "black",
+                    ]
+                    + expected_additional_arguments_fn(rev.path),
+                    cwd=cwd,
+                )
+            ],
+        )
+
+    @testing.variation("config", ["ini", "toml"])
+    def test_module(self, config):
+        if config.ini:
+            input_config = """
+[post_write_hooks]
+hooks = black
+black.type = module
+black.module = black
+black.options = -l 79
+            """
+        else:
+            input_config = """
+    [[tool.alembic.post_write_hooks]]
+    name = "black"
+    type = "module"
+    module = "black"
+    options = "-l 79"
+    """
+
+        def expected_additional_arguments_fn(rev_path):
+            return [rev_path, "-l", "79"]
+
+        self._run_black_module_with_config(
+            input_config,
+            expected_additional_arguments_fn,
+            use_toml=config.toml,
+        )
+
+    @combinations(True, False, argnames="posix")
+    @testing.variation("config", ["ini", "toml"])
+    def test_module_filename_interpolation(self, posix, config):
+        if config.ini:
+            input_config = """
+[post_write_hooks]
+hooks = black
+black.type = module
+black.module = black
+black.options = arg1 REVISION_SCRIPT_FILENAME 'multi-word arg' \
+    --flag1='REVISION_SCRIPT_FILENAME'
+        """
+        else:
+            input_config = """
+[[tool.alembic.post_write_hooks]]
+name = "black"
+type = "module"
+module = "black"
+options = "arg1 REVISION_SCRIPT_FILENAME 'multi-word arg' --flag1='REVISION_SCRIPT_FILENAME'"
+"""  # noqa: E501
+
+        def expected_additional_arguments_fn(rev_path):
+            if compat.is_posix:
+                return [
+                    "arg1",
+                    rev_path,
+                    "multi-word arg",
+                    "--flag1=" + rev_path,
+                ]
+            else:
+                return [
+                    "arg1",
+                    rev_path,
+                    "'multi-word arg'",
+                    "--flag1='%s'" % rev_path,
+                ]
+
+        with mock.patch("alembic.util.compat.is_posix", posix):
+            self._run_black_module_with_config(
+                input_config,
+                expected_additional_arguments_fn,
+                use_toml=config.toml,
+            )
+
+    def test_module_path_in_config(self):
+        input_config = """
+[post_write_hooks]
+hooks = black
+black.type = module
+black.module = black
+black.options = arg1 REVISION_SCRIPT_FILENAME --config %(here)s/pyproject.toml
+        """
+
+        def expected_additional_arguments_fn(rev_path):
+            return [
+                "arg1",
+                rev_path,
+                "--config",
+                Path(_get_staging_directory(), "pyproject.toml")
+                .absolute()
+                .as_posix(),
+            ]
+
+        self._run_black_module_with_config(
+            input_config, expected_additional_arguments_fn
+        )
+
+    def test_module_black_with_cwd(self):
+        input_config = """
+[post_write_hooks]
+hooks = black
+black.type = module
+black.module = black
+black.cwd = /path/to/cwd
+        """
+
+        def expected_additional_arguments_fn(rev_path):
+            return [rev_path]
+
+        self._run_black_module_with_config(
+            input_config, expected_additional_arguments_fn, cwd="/path/to/cwd"
+        )