From c746e1b6eb0814e16c737244727f3b30d99f1a15 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Tue, 29 Apr 2025 15:49:35 -0400 Subject: [PATCH] Refactor `CommandLine` command registration Added new :meth:`.CommandLine.register_command` method to :class:`.CommandLine`, intended to facilitate adding custom commands to Alembic's command line tool with minimal code required; previously this logic was embedded internally and was not publicly accessible. A new recipe demonstrating this use is added. Pull request courtesy Mikhail Bulash. Fixes: #1610 Closes: #1611 Pull-request: https://github.com/sqlalchemy/alembic/pull/1611 Pull-request-sha: d2ffce008317508449de9bd0bce3ff9076c5d90f Change-Id: Icab9978357915e81bb52a9f0717f21d5ee7f8341 --- alembic/config.py | 434 ++++++++++++++++++--------------- docs/build/cookbook.rst | 60 +++++ docs/build/unreleased/1610.rst | 14 ++ tests/test_config.py | 26 ++ 4 files changed, 331 insertions(+), 203 deletions(-) create mode 100644 docs/build/unreleased/1610.rst diff --git a/alembic/config.py b/alembic/config.py index 2c52e7cd..18ab0f16 100644 --- a/alembic/config.py +++ b/alembic/config.py @@ -12,6 +12,7 @@ from typing import Dict from typing import Mapping from typing import Optional from typing import overload +from typing import Protocol from typing import Sequence from typing import TextIO from typing import Union @@ -344,176 +345,171 @@ class MessagingOptions(TypedDict, total=False): quiet: bool +class CommandFunction(Protocol): + """A function that may be registered in the CLI as an alembic command. + It must be a named function and it must accept a :class:`.Config` object + as the first argument. + + .. versionadded:: 1.15.3 + + """ + + __name__: str + + def __call__(self, config: Config, *args: Any, **kwargs: Any) -> Any: ... + + class CommandLine: + """Provides the command line interface to Alembic.""" + def __init__(self, prog: Optional[str] = None) -> None: self._generate_args(prog) - def _generate_args(self, prog: Optional[str]) -> None: - def add_options( - fn: Any, parser: Any, positional: Any, kwargs: Any - ) -> None: - kwargs_opts = { - "template": ( - "-t", - "--template", - dict( - default="generic", - type=str, - help="Setup template for use with 'init'", - ), - ), - "message": ( - "-m", - "--message", - dict( - type=str, help="Message string to use with 'revision'" - ), - ), - "sql": ( - "--sql", - dict( - action="store_true", - help="Don't emit SQL to database - dump to " - "standard output/file instead. See docs on " - "offline mode.", - ), - ), - "tag": ( - "--tag", - dict( - type=str, - help="Arbitrary 'tag' name - can be used by " - "custom env.py scripts.", - ), - ), - "head": ( - "--head", - dict( - type=str, - help="Specify head revision or @head " - "to base new revision on.", - ), - ), - "splice": ( - "--splice", - dict( - action="store_true", - help="Allow a non-head revision as the " - "'head' to splice onto", - ), - ), - "depends_on": ( - "--depends-on", - dict( - action="append", - help="Specify one or more revision identifiers " - "which this revision should depend on.", - ), - ), - "rev_id": ( - "--rev-id", - dict( - type=str, - help="Specify a hardcoded revision id instead of " - "generating one", - ), - ), - "version_path": ( - "--version-path", - dict( - type=str, - help="Specify specific path from config for " - "version file", - ), - ), - "branch_label": ( - "--branch-label", - dict( - type=str, - help="Specify a branch label to apply to the " - "new revision", - ), - ), - "verbose": ( - "-v", - "--verbose", - dict(action="store_true", help="Use more verbose output"), - ), - "resolve_dependencies": ( - "--resolve-dependencies", - dict( - action="store_true", - help="Treat dependency versions as down revisions", - ), - ), - "autogenerate": ( - "--autogenerate", - dict( - action="store_true", - help="Populate revision script with candidate " - "migration operations, based on comparison " - "of database to model.", - ), - ), - "rev_range": ( - "-r", - "--rev-range", - dict( - action="store", - help="Specify a revision range; " - "format is [start]:[end]", - ), - ), - "indicate_current": ( - "-i", - "--indicate-current", - dict( - action="store_true", - help="Indicate the current revision", - ), - ), - "purge": ( - "--purge", - dict( - action="store_true", - help="Unconditionally erase the version table " - "before stamping", - ), - ), - "package": ( - "--package", - dict( - action="store_true", - help="Write empty __init__.py files to the " - "environment and version locations", - ), - ), - } - positional_help = { - "directory": "location of scripts directory", - "revision": "revision identifier", - "revisions": "one or more revisions, or 'heads' for all heads", - } - for arg in kwargs: - if arg in kwargs_opts: - args = kwargs_opts[arg] - args, kw = args[0:-1], args[-1] - parser.add_argument(*args, **kw) - - for arg in positional: - if ( - arg == "revisions" - or fn in positional_translations - and positional_translations[fn][arg] == "revisions" - ): - subparser.add_argument( - "revisions", - nargs="+", - help=positional_help.get("revisions"), - ) - else: - subparser.add_argument(arg, help=positional_help.get(arg)) + _KWARGS_OPTS = { + "template": ( + "-t", + "--template", + dict( + default="generic", + type=str, + help="Setup template for use with 'init'", + ), + ), + "message": ( + "-m", + "--message", + dict(type=str, help="Message string to use with 'revision'"), + ), + "sql": ( + "--sql", + dict( + action="store_true", + help="Don't emit SQL to database - dump to " + "standard output/file instead. See docs on " + "offline mode.", + ), + ), + "tag": ( + "--tag", + dict( + type=str, + help="Arbitrary 'tag' name - can be used by " + "custom env.py scripts.", + ), + ), + "head": ( + "--head", + dict( + type=str, + help="Specify head revision or @head " + "to base new revision on.", + ), + ), + "splice": ( + "--splice", + dict( + action="store_true", + help="Allow a non-head revision as the 'head' to splice onto", + ), + ), + "depends_on": ( + "--depends-on", + dict( + action="append", + help="Specify one or more revision identifiers " + "which this revision should depend on.", + ), + ), + "rev_id": ( + "--rev-id", + dict( + type=str, + help="Specify a hardcoded revision id instead of " + "generating one", + ), + ), + "version_path": ( + "--version-path", + dict( + type=str, + help="Specify specific path from config for version file", + ), + ), + "branch_label": ( + "--branch-label", + dict( + type=str, + help="Specify a branch label to apply to the new revision", + ), + ), + "verbose": ( + "-v", + "--verbose", + dict(action="store_true", help="Use more verbose output"), + ), + "resolve_dependencies": ( + "--resolve-dependencies", + dict( + action="store_true", + help="Treat dependency versions as down revisions", + ), + ), + "autogenerate": ( + "--autogenerate", + dict( + action="store_true", + help="Populate revision script with candidate " + "migration operations, based on comparison " + "of database to model.", + ), + ), + "rev_range": ( + "-r", + "--rev-range", + dict( + action="store", + help="Specify a revision range; format is [start]:[end]", + ), + ), + "indicate_current": ( + "-i", + "--indicate-current", + dict( + action="store_true", + help="Indicate the current revision", + ), + ), + "purge": ( + "--purge", + dict( + action="store_true", + help="Unconditionally erase the version table before stamping", + ), + ), + "package": ( + "--package", + dict( + action="store_true", + help="Write empty __init__.py files to the " + "environment and version locations", + ), + ), + } + _POSITIONAL_OPTS = { + "directory": dict(help="location of scripts directory"), + "revision": dict( + help="revision identifier", + ), + "revisions": dict( + nargs="+", + help="one or more revisions, or 'heads' for all heads", + ), + } + _POSITIONAL_TRANSLATIONS: dict[Any, dict[str, str]] = { + command.stamp: {"revision": "revisions"} + } + def _generate_args(self, prog: Optional[str]) -> None: parser = ArgumentParser(prog=prog) parser.add_argument( @@ -532,7 +528,7 @@ class CommandLine: "--name", type=str, default="alembic", - help="Name of section in .ini file to " "use for Alembic config", + help="Name of section in .ini file to use for Alembic config", ) parser.add_argument( "-x", @@ -552,50 +548,81 @@ class CommandLine: action="store_true", help="Do not log to std output.", ) - subparsers = parser.add_subparsers() - - positional_translations: Dict[Any, Any] = { - command.stamp: {"revision": "revisions"} - } - for fn in [getattr(command, n) for n in dir(command)]: + self.subparsers = parser.add_subparsers() + alembic_commands = ( + cast(CommandFunction, fn) + for fn in (getattr(command, name) for name in dir(command)) if ( inspect.isfunction(fn) and fn.__name__[0] != "_" and fn.__module__ == "alembic.command" - ): - spec = compat.inspect_getfullargspec(fn) - if spec[3] is not None: - positional = spec[0][1 : -len(spec[3])] - kwarg = spec[0][-len(spec[3]) :] - else: - positional = spec[0][1:] - kwarg = [] - - if fn in positional_translations: - positional = [ - positional_translations[fn].get(name, name) - for name in positional - ] - - # parse first line(s) of helptext without a line break - help_ = fn.__doc__ - if help_: - help_text = [] - for line in help_.split("\n"): - if not line.strip(): - break - else: - help_text.append(line.strip()) - else: - help_text = [] - subparser = subparsers.add_parser( - fn.__name__, help=" ".join(help_text) - ) - add_options(fn, subparser, positional, kwarg) - subparser.set_defaults(cmd=(fn, positional, kwarg)) + ) + ) + + for fn in alembic_commands: + self.register_command(fn) + self.parser = parser + def register_command(self, fn: CommandFunction) -> None: + """Registers a function as a CLI subcommand. The subcommand name + matches the function name, the arguments are extracted from the + signature and the help text is read from the docstring. + + .. versionadded:: 1.15.3 + + .. seealso:: + + :ref:`custom_commandline` + """ + + positional, kwarg, help_text = self._inspect_function(fn) + + subparser = self.subparsers.add_parser(fn.__name__, help=help_text) + subparser.set_defaults(cmd=(fn, positional, kwarg)) + + for arg in kwarg: + if arg in self._KWARGS_OPTS: + kwarg_opt = self._KWARGS_OPTS[arg] + args, opts = kwarg_opt[0:-1], kwarg_opt[-1] + subparser.add_argument(*args, **opts) # type:ignore + + for arg in positional: + opts = self._POSITIONAL_OPTS.get(arg, {}) + subparser.add_argument(arg, **opts) # type:ignore + + def _inspect_function(self, fn: CommandFunction) -> tuple[Any, Any, str]: + spec = compat.inspect_getfullargspec(fn) + if spec[3] is not None: + positional = spec[0][1 : -len(spec[3])] + kwarg = spec[0][-len(spec[3]) :] + else: + positional = spec[0][1:] + kwarg = [] + + if fn in self._POSITIONAL_TRANSLATIONS: + positional = [ + self._POSITIONAL_TRANSLATIONS[fn].get(name, name) + for name in positional + ] + + # parse first line(s) of helptext without a line break + help_ = fn.__doc__ + if help_: + help_lines = [] + for line in help_.split("\n"): + if not line.strip(): + break + else: + help_lines.append(line.strip()) + else: + help_lines = [] + + help_text = " ".join(help_lines) + + return positional, kwarg, help_text + def run_cmd(self, config: Config, options: Namespace) -> None: fn, positional, kwarg = options.cmd @@ -612,6 +639,7 @@ class CommandLine: util.err(str(e), **config.messaging_opts) def main(self, argv: Optional[Sequence[str]] = None) -> None: + """Executes the command line with the provided arguments.""" options = self.parser.parse_args(argv) if not hasattr(options, "cmd"): # see http://bugs.python.org/issue9253, argparse diff --git a/docs/build/cookbook.rst b/docs/build/cookbook.rst index ce5fb543..e52753a8 100644 --- a/docs/build/cookbook.rst +++ b/docs/build/cookbook.rst @@ -1615,3 +1615,63 @@ The application maintains a version of schema with both versions. Writes are performed on both places, while the background script move all the remaining data across. This technique is very challenging and time demanding, since it requires custom application logic to handle the intermediate states. + +.. _custom_commandline: + +Extend ``CommandLine`` with custom commands +=========================================== + +.. versionadded:: 1.15.3 + +While Alembic does not have a plugin system that would allow transparently extending the original ``alembic`` CLI +with additional commands, it is possible to create your own instance of :class:`.CommandLine` and extend that via +:meth:`.CommandLine.register_command`. + +.. code-block:: python + + # myalembic.py + + from alembic.config import CommandLine, Config + + + def frobnicate(config: Config, revision: str) -> None: + """Frobnicates according to the frobnication specification. + + :param config: a :class:`.Config` instance + :param revision: the revision to frobnicate + """ + + config.print_stdout(f"Revision {revision} successfully frobnicated.") + + + def main(): + cli = CommandLine() + cli.register_command(frobnicate) + cli.main() + + if __name__ == "__main__": + main() + +Any named function may be registered as a command, provided it accepts a :class:`.Config` object as the first argument; +a docstring is also recommended as it will show up in the help output of the CLI. + +.. code-block:: + + $ python -m myalembic -h + ... + positional arguments: + {branches,check,current,downgrade,edit,ensure_version,heads,history,init,list_templates,merge,revision,show,stamp,upgrade,frobnicate} + ... + frobnicate Frobnicates according to the frobnication specification. + + $ python -m myalembic frobnicate -h + usage: myalembic.py frobnicate [-h] revision + + positional arguments: + revision revision identifier + + optional arguments: + -h, --help show this help message and exit + + $ python -m myalembic frobnicate 42 + Revision 42 successfully frobnicated. diff --git a/docs/build/unreleased/1610.rst b/docs/build/unreleased/1610.rst new file mode 100644 index 00000000..b6f37d09 --- /dev/null +++ b/docs/build/unreleased/1610.rst @@ -0,0 +1,14 @@ +.. change:: + :tags: usecase, commands + :tickets: 1610 + + Added new :meth:`.CommandLine.register_command` method to + :class:`.CommandLine`, intended to facilitate adding custom commands to + Alembic's command line tool with minimal code required; previously this + logic was embedded internally and was not publicly accessible. A new + recipe demonstrating this use is added. Pull request courtesy Mikhail + Bulash. + + .. seealso:: + + :ref:`custom_commandline` diff --git a/tests/test_config.py b/tests/test_config.py index a98994c4..0fad0dda 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -256,3 +256,29 @@ class TemplateOutputEncodingTest(TestBase): self.cfg.set_main_option("output_encoding", "latin-1") script = ScriptDirectory.from_config(self.cfg) eq_(script.output_encoding, "latin-1") + + +class CommandLineTest(TestBase): + def test_register_command(self): + cli = config.CommandLine() + + fake_stdout = [] + + def frobnicate(config: config.Config, revision: str) -> None: + """Frobnicates the revision. + + :param config: a :class:`.Config` instance + :param revision: the revision to frobnicate + """ + + fake_stdout.append(f"Revision {revision} frobnicated.") + + cli.register_command(frobnicate) + + help_text = cli.parser.format_help() + assert frobnicate.__name__ in help_text + assert frobnicate.__doc__.split("\n")[0] in help_text + + cli.main(["frobnicate", "abc42"]) + + assert fake_stdout == ["Revision abc42 frobnicated."] -- 2.47.3