Public interface, docs and a test

diff --git a/alembic/config.py b/alembic/config.py
index 1866909fc194fe6480a6b9571e8e72fb48a1238c..1ff2e4623e8e4e3b609092ff0b40c7697da4a7fe 100644 (file)
--- 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 ce5fb5438e8ced95aef47fc76ab5794359bd223f..d8ee19b6d2a2b8e966358e8794a43eee990c5394 100644 (file)
--- 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 a98994c4d22914104f5fb90a8eb86eb468d62567..6a69463cd8f4358749dc3bdba0162fd132a48be8 100644 (file)
--- 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."
+        ]