From: Mikhail Bulash Date: Thu, 24 Apr 2025 16:17:57 +0000 (+0200) Subject: Public interface, docs and a test X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=2e0df014d83e5f3fc49348d6c0fc67e8cc21d164;p=thirdparty%2Fsqlalchemy%2Falembic.git Public interface, docs and a test --- 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 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,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 @@ class CommandLine: 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 @@ class CommandLine: ) 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 @@ class CommandLine: 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 @@ 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..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 @@ 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 == [ + f"Revision abc42 frobnicated." + ]