From e1d64546e3e7d24a6723065ba54760766fa758e4 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Thu, 20 Feb 2025 20:57:09 +0100 Subject: [PATCH 1/8] Refactor `CommandLine` command registration Fixes: #1610 --- alembic/config.py | 405 +++++++++++++++++++++++----------------------- 1 file changed, 203 insertions(+), 202 deletions(-) diff --git a/alembic/config.py b/alembic/config.py index 2c52e7cd..ecf793e6 100644 --- a/alembic/config.py +++ b/alembic/config.py @@ -348,172 +348,144 @@ class CommandLine: 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_HELP = { + "directory": "location of scripts directory", + "revision": "revision identifier", + "revisions": "one or more revisions, or 'heads' for all heads", + } + _POSITIONAL_TRANSLATIONS: dict[str, str] = { + command.stamp: {"revision": "revisions"} + } + def _generate_args(self, prog: Optional[str]) -> None: parser = ArgumentParser(prog=prog) parser.add_argument( @@ -532,7 +504,7 @@ def add_options( "--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,49 +524,78 @@ def add_options( 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 = ( + 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" + ) + ) + + for fn in alembic_commands: + self._register_command(fn) + + self.parser = parser + + def _register_command(self, fn: Any) -> None: + fn, positional, kwarg, help_text = self._parse_command_from_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: + args = self._KWARGS_OPTS[arg] + args, kw = args[0:-1], args[-1] + subparser.add_argument(*args, **kw) + + for arg in positional: + if ( + arg == "revisions" + or fn in self._POSITIONAL_TRANSLATIONS + and self._POSITIONAL_TRANSLATIONS[fn][arg] == "revisions" ): - 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) + subparser.add_argument( + "revisions", + nargs="+", + help=self._POSITIONAL_HELP.get("revisions"), ) - add_options(fn, subparser, positional, kwarg) - subparser.set_defaults(cmd=(fn, positional, kwarg)) - self.parser = parser + else: + subparser.add_argument(arg, help=self._POSITIONAL_HELP.get(arg)) + + def _parse_command_from_function(self, fn: Any) -> tuple[Any, 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 fn, positional, kwarg, help_text def run_cmd(self, config: Config, options: Namespace) -> None: fn, positional, kwarg = options.cmd From aebe80a0c826c7e2e91b90dc542dc499e537acf8 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Thu, 27 Feb 2025 15:09:15 +0100 Subject: [PATCH 2/8] fixup! Refactor `CommandLine` command registration --- alembic/config.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/alembic/config.py b/alembic/config.py index ecf793e6..73921ec6 100644 --- a/alembic/config.py +++ b/alembic/config.py @@ -481,7 +481,7 @@ def __init__(self, prog: Optional[str] = None) -> None: "revision": "revision identifier", "revisions": "one or more revisions, or 'heads' for all heads", } - _POSITIONAL_TRANSLATIONS: dict[str, str] = { + _POSITIONAL_TRANSLATIONS: dict[Any, dict[str, str]] = { command.stamp: {"revision": "revisions"} } @@ -549,9 +549,9 @@ def _register_command(self, fn: Any) -> None: for arg in kwarg: if arg in self._KWARGS_OPTS: - args = self._KWARGS_OPTS[arg] - args, kw = args[0:-1], args[-1] - subparser.add_argument(*args, **kw) + kwarg_opt = self._KWARGS_OPTS[arg] + args, kw = kwarg_opt[0:-1], kwarg_opt[-1] + subparser.add_argument(*args, **kw) # type:ignore for arg in positional: if ( From 369059eb1e7d569575407eb5ab191c5728e7b081 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Thu, 27 Feb 2025 15:15:41 +0100 Subject: [PATCH 3/8] fixup! Refactor `CommandLine` command registration --- alembic/util/compat.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/alembic/util/compat.py b/alembic/util/compat.py index fa8bc02b..08ebe25f 100644 --- a/alembic/util/compat.py +++ b/alembic/util/compat.py @@ -7,6 +7,7 @@ import os import sys import typing +from typing import cast from typing import Any from typing import List from typing import Optional @@ -55,9 +56,9 @@ def close(self) -> None: def importlib_metadata_get(group: str) -> Sequence[EntryPoint]: ep = importlib_metadata.entry_points() if hasattr(ep, "select"): - return ep.select(group=group) + return cast(Sequence[EntryPoint], ep.select(group=group)) else: - return ep.get(group, ()) # type: ignore + return ep.get(group, ()) def formatannotation_fwdref( From ed96df233230d362f95b5a4ab6f2c92175784b06 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Thu, 6 Mar 2025 12:56:15 +0100 Subject: [PATCH 4/8] fixup! Refactor `CommandLine` command registration --- alembic/config.py | 19 ++++++++++++------- alembic/util/compat.py | 6 +++--- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/alembic/config.py b/alembic/config.py index 73921ec6..f2046703 100644 --- a/alembic/config.py +++ b/alembic/config.py @@ -376,7 +376,8 @@ def __init__(self, prog: Optional[str] = None) -> None: "--tag", dict( type=str, - help="Arbitrary 'tag' name - can be used by custom env.py scripts.", + help="Arbitrary 'tag' name - can be used by " + "custom env.py scripts.", ), ), "head": ( @@ -406,7 +407,8 @@ def __init__(self, prog: Optional[str] = None) -> None: "--rev-id", dict( type=str, - help="Specify a hardcoded revision id instead of generating one", + help="Specify a hardcoded revision id instead of " + "generating one", ), ), "version_path": ( @@ -542,7 +544,7 @@ def _generate_args(self, prog: Optional[str]) -> None: self.parser = parser def _register_command(self, fn: Any) -> None: - fn, positional, kwarg, help_text = self._parse_command_from_function(fn) + 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)) @@ -565,9 +567,11 @@ def _register_command(self, fn: Any) -> None: help=self._POSITIONAL_HELP.get("revisions"), ) else: - subparser.add_argument(arg, help=self._POSITIONAL_HELP.get(arg)) + subparser.add_argument( + arg, help=self._POSITIONAL_HELP.get(arg) + ) - def _parse_command_from_function(self, fn: Any) -> tuple[Any, Any, Any, str]: + def _inspect_function(self, fn: Any) -> tuple[Any, Any, str]: spec = compat.inspect_getfullargspec(fn) if spec[3] is not None: positional = spec[0][1 : -len(spec[3])] @@ -578,7 +582,8 @@ def _parse_command_from_function(self, fn: Any) -> tuple[Any, Any, Any, str]: if fn in self._POSITIONAL_TRANSLATIONS: positional = [ - self._POSITIONAL_TRANSLATIONS[fn].get(name, name) for name in positional + self._POSITIONAL_TRANSLATIONS[fn].get(name, name) + for name in positional ] # parse first line(s) of helptext without a line break @@ -595,7 +600,7 @@ def _parse_command_from_function(self, fn: Any) -> tuple[Any, Any, Any, str]: help_text = " ".join(help_lines) - return fn, positional, kwarg, help_text + return positional, kwarg, help_text def run_cmd(self, config: Config, options: Namespace) -> None: fn, positional, kwarg = options.cmd diff --git a/alembic/util/compat.py b/alembic/util/compat.py index 08ebe25f..6c4e3d40 100644 --- a/alembic/util/compat.py +++ b/alembic/util/compat.py @@ -7,8 +7,8 @@ import os import sys import typing -from typing import cast from typing import Any +from typing import cast from typing import List from typing import Optional from typing import Sequence @@ -56,9 +56,9 @@ def close(self) -> None: def importlib_metadata_get(group: str) -> Sequence[EntryPoint]: ep = importlib_metadata.entry_points() if hasattr(ep, "select"): - return cast(Sequence[EntryPoint], ep.select(group=group)) + return ep.select(group=group) else: - return ep.get(group, ()) + return ep.get(group, ()) # type: ignore def formatannotation_fwdref( From 16f8e85fce1bd213916e6cef7e5bbee220d52786 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Thu, 6 Mar 2025 14:30:04 +0100 Subject: [PATCH 5/8] Simplify _POSITIONAL_OPTS --- alembic/config.py | 34 ++++++++++++++-------------------- 1 file changed, 14 insertions(+), 20 deletions(-) diff --git a/alembic/config.py b/alembic/config.py index f2046703..1866909f 100644 --- a/alembic/config.py +++ b/alembic/config.py @@ -478,10 +478,15 @@ def __init__(self, prog: Optional[str] = None) -> None: ), ), } - _POSITIONAL_HELP = { - "directory": "location of scripts directory", - "revision": "revision identifier", - "revisions": "one or more revisions, or 'heads' for all heads", + _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"} @@ -552,24 +557,13 @@ def _register_command(self, fn: Any) -> None: for arg in kwarg: if arg in self._KWARGS_OPTS: kwarg_opt = self._KWARGS_OPTS[arg] - args, kw = kwarg_opt[0:-1], kwarg_opt[-1] - subparser.add_argument(*args, **kw) # type:ignore + args, opts = kwarg_opt[0:-1], kwarg_opt[-1] + subparser.add_argument(*args, **opts) # type:ignore for arg in positional: - if ( - arg == "revisions" - or fn in self._POSITIONAL_TRANSLATIONS - and self._POSITIONAL_TRANSLATIONS[fn][arg] == "revisions" - ): - subparser.add_argument( - "revisions", - nargs="+", - help=self._POSITIONAL_HELP.get("revisions"), - ) - else: - subparser.add_argument( - arg, help=self._POSITIONAL_HELP.get(arg) - ) + if arg in self._POSITIONAL_OPTS: + opts = self._POSITIONAL_OPTS[arg] + subparser.add_argument(arg, **opts) # type:ignore def _inspect_function(self, fn: Any) -> tuple[Any, Any, str]: spec = compat.inspect_getfullargspec(fn) From 300693537df591a06c593e4f1032e0e581133e62 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Thu, 6 Mar 2025 14:39:56 +0100 Subject: [PATCH 6/8] fixup! Refactor `CommandLine` command registration --- alembic/util/compat.py | 1 - 1 file changed, 1 deletion(-) diff --git a/alembic/util/compat.py b/alembic/util/compat.py index 6c4e3d40..fa8bc02b 100644 --- a/alembic/util/compat.py +++ b/alembic/util/compat.py @@ -8,7 +8,6 @@ import sys import typing from typing import Any -from typing import cast from typing import List from typing import Optional from typing import Sequence From 2e0df014d83e5f3fc49348d6c0fc67e8cc21d164 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Thu, 24 Apr 2025 18:17:57 +0200 Subject: [PATCH 7/8] Public interface, docs and a test --- alembic/config.py | 31 +++++++++++++++++++--- docs/build/cookbook.rst | 58 +++++++++++++++++++++++++++++++++++++++++ tests/test_config.py | 28 ++++++++++++++++++++ 3 files changed, 113 insertions(+), 4 deletions(-) diff --git a/alembic/config.py b/alembic/config.py index 1866909f..1ff2e462 100644 --- a/alembic/config.py +++ b/alembic/config.py @@ -12,6 +12,7 @@ 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,7 +345,19 @@ 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. + """ + + __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) @@ -534,7 +547,7 @@ def _generate_args(self, prog: Optional[str]) -> None: self.subparsers = parser.add_subparsers() alembic_commands = ( - fn + cast(CommandFunction, fn) for fn in (getattr(command, name) for name in dir(command)) if ( inspect.isfunction(fn) @@ -544,11 +557,20 @@ def _generate_args(self, prog: Optional[str]) -> None: ) for fn in alembic_commands: - self._register_command(fn) + self.register_command(fn) self.parser = parser - def _register_command(self, fn: Any) -> None: + 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. + + .. seealso:: + + :ref:`custom_commandline` + """ + positional, kwarg, help_text = self._inspect_function(fn) subparser = self.subparsers.add_parser(fn.__name__, help=help_text) @@ -565,7 +587,7 @@ def _register_command(self, fn: Any) -> None: opts = self._POSITIONAL_OPTS[arg] subparser.add_argument(arg, **opts) # type:ignore - def _inspect_function(self, fn: Any) -> tuple[Any, Any, str]: + 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])] @@ -612,6 +634,7 @@ def run_cmd(self, config: Config, options: Namespace) -> None: 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..d8ee19b6 100644 --- a/docs/build/cookbook.rst +++ b/docs/build/cookbook.rst @@ -1615,3 +1615,61 @@ 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 +=========================================== + +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/tests/test_config.py b/tests/test_config.py index a98994c4..6a69463c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -256,3 +256,31 @@ def test_setting(self): 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 == [ + f"Revision abc42 frobnicated." + ] From 2d4e7db96278d926a5d53bb7e4974853fb12c306 Mon Sep 17 00:00:00 2001 From: Mikhail Bulash Date: Thu, 24 Apr 2025 18:54:11 +0200 Subject: [PATCH 8/8] fixup! Public interface, docs and a test --- alembic/config.py | 9 +++++---- tests/test_config.py | 6 ++---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/alembic/config.py b/alembic/config.py index 1ff2e462..47fdb3f8 100644 --- a/alembic/config.py +++ b/alembic/config.py @@ -346,8 +346,9 @@ class MessagingOptions(TypedDict, total=False): 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. + """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. """ __name__: str @@ -563,8 +564,8 @@ def _generate_args(self, prog: Optional[str]) -> None: 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. + matches the function name, the arguments are extracted from the + signature and the help text is read from the docstring. .. seealso:: diff --git a/tests/test_config.py b/tests/test_config.py index 6a69463c..0fad0dda 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -263,7 +263,7 @@ def test_register_command(self): cli = config.CommandLine() fake_stdout = [] - + def frobnicate(config: config.Config, revision: str) -> None: """Frobnicates the revision. @@ -281,6 +281,4 @@ def frobnicate(config: config.Config, revision: str) -> None: cli.main(["frobnicate", "abc42"]) - assert fake_stdout == [ - f"Revision abc42 frobnicated." - ] + assert fake_stdout == ["Revision abc42 frobnicated."]