Factor out common hook functionality

diff --git a/alembic/script/write_hooks.py b/alembic/script/write_hooks.py
index cf96e91d8af7b499dc9deaa2789923012e5c2d27..f40bb35f6a8c8878fe4fe4181f235e9fc010b36b 100644 (file)
--- a/alembic/script/write_hooks.py
+++ b/alembic/script/write_hooks.py
@@ -113,25 +113,19 @@ def _parse_cmdline_options(cmdline_options_str: str, path: str) -> List[str]:
     return cmdline_options_list
 
 
-@register("console_scripts")
-def console_scripts(
-    path: str, options: dict, ignore_output: bool = False
-) -> None:
+def _get_required_option(options: dict, name: str) -> str:
     try:
-        entrypoint_name = options["entrypoint"]
+        return options[name]
     except KeyError as ke:
         raise util.CommandError(
-            f"Key {options['_hook_name']}.entrypoint is required for post "
+            f"Key {options['_hook_name']}.{name} is required for post "
             f"write hook {options['_hook_name']!r}"
         ) from ke
-    for entry in compat.importlib_metadata_get("console_scripts"):
-        if entry.name == entrypoint_name:
-            impl: Any = entry
-            break
-    else:
-        raise util.CommandError(
-            f"Could not find entrypoint console_scripts.{entrypoint_name}"
-        )
+
+
+def _run_hook(
+    path: str, options: dict, ignore_output: bool, command: List[str]
+) -> None:
     cwd: Optional[str] = options.get("cwd", None)
     cmdline_options_str = options.get("options", "")
     cmdline_options_list = _parse_cmdline_options(cmdline_options_str, path)
@@ -140,73 +134,43 @@ def console_scripts(
     if ignore_output:
         kw["stdout"] = kw["stderr"] = subprocess.DEVNULL
 
-    subprocess.run(
-        [
-            sys.executable,
-            "-c",
-            f"import {impl.module}; {impl.module}.{impl.attr}()",
-        ]
-        + cmdline_options_list,
-        cwd=cwd,
-        **kw,
-    )
+    subprocess.run([*command, *cmdline_options_list], cwd=cwd, **kw)
 
 
-@register("exec")
-def exec_(path: str, options: dict, ignore_output: bool = False) -> None:
-    try:
-        executable = options["executable"]
-    except KeyError as ke:
+@register("console_scripts")
+def console_scripts(
+    path: str, options: dict, ignore_output: bool = False
+) -> None:
+    entrypoint_name = _get_required_option(options, "entrypoint")
+    for entry in compat.importlib_metadata_get("console_scripts"):
+        if entry.name == entrypoint_name:
+            impl: Any = entry
+            break
+    else:
         raise util.CommandError(
-            f"Key {options['_hook_name']}.executable is required for post "
-            f"write hook {options['_hook_name']!r}"
-        ) from ke
-    cwd: Optional[str] = options.get("cwd", None)
-    cmdline_options_str = options.get("options", "")
-    cmdline_options_list = _parse_cmdline_options(cmdline_options_str, path)
+            f"Could not find entrypoint console_scripts.{entrypoint_name}"
+        )
 
-    kw: Dict[str, Any] = {}
-    if ignore_output:
-        kw["stdout"] = kw["stderr"] = subprocess.DEVNULL
+    command = [
+        sys.executable,
+        "-c",
+        f"import {impl.module}; {impl.module}.{impl.attr}()",
+    ]
+    _run_hook(path, options, ignore_output, command)
 
-    subprocess.run(
-        [
-            executable,
-            *cmdline_options_list,
-        ],
-        cwd=cwd,
-        **kw,
-    )
+
+@register("exec")
+def exec_(path: str, options: dict, ignore_output: bool = False) -> None:
+    executable = _get_required_option(options, "executable")
+    _run_hook(path, options, ignore_output, command=[executable])
 
 
 @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
+    module_name = _get_required_option(options, "module")
 
     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,
-    )
+    command = [sys.executable, "-m", module_name]
+    _run_hook(path, options, ignore_output, command)