Add mechanism for extending existing spin commands

Closes #242

Author: stefanv

README.md

@@ -201,52 +201,37 @@ click options or function keywords:
 
 ### Advanced: adding arguments to built-in commands
 
-Instead of rewriting a command from scratch, a project may want to add a flag to a built-in `spin` command, or perhaps do some pre- or post-processing.
-For this, we have to use an internal Click concept called a [context](https://click.palletsprojects.com/en/8.1.x/complex/#contexts).
-Fortunately, we don't need to know anything about contexts other than that they allow us to execute commands within commands.
+Instead of rewriting a command from scratch, a project may simply want to add a flag to an existing `spin` command, or perhaps do some pre- or post-processing.
+For this purpose, we provide the `spin.util.extend_cmd` decorator.
 
-We proceed by duplicating the function header of the existing command, and adding our own flag:
+Here, we show how to add a `--extra` flag to the existing `build` function:
 
 ```python
 from spin.cmds import meson
 
-# Take this from the built-in implementation, in `spin.cmds.meson.build`:
 
+@click.option("-e", "--extra", help="Extra test flag")
+@util.extend_command(spin.cmds.meson.build)
+def build_extend(*, parent_callback, extra=None, **kwargs):
+    """
+    This version of build also provides the EXTRA flag, that can be used
+    to specify an extra integer argument.
+    """
+    print(f"Preparing for build with {extra=}")
+    parent_callback(**kwargs)
+    print("Finalizing build...")
+```
 
-@click.command()
-@click.argument("meson_args", nargs=-1)
-@click.option("-j", "--jobs", help="Number of parallel tasks to launch", type=int)
-@click.option("--clean", is_flag=True, help="Clean build directory before build")
-@click.option(
-    "-v", "--verbose", is_flag=True, help="Print all build output, even installation"
-)
-
-# This is our new option
-@click.option("--custom-arg/--no-custom-arg")
-
-# This tells spin that we will need a context, which we
-# can use to invoke the built-in command
-@click.pass_context
-
-# This is the original function signature, plus our new flag
-def build(ctx, meson_args, jobs=None, clean=False, verbose=False, custom_arg=False):
-    """Docstring goes here. You may want to copy and customize the original."""
-
-    # Do something with the new option
-    print("The value of custom arg is:", custom_arg)
-
-    # The spin `build` command doesn't know anything about `custom_arg`,
-    # so don't send it on.
-    del ctx.params["custom_arg"]
+Note that `build_extend` receives the parent command callback (the function the `build` command would have executed) as its first argument.
 
-    # Call the built-in `build` command, passing along
-    # all arguments and options.
-    ctx.forward(meson.build)
+The matching entry in `pyproject.toml` is:
 
-    # Also see:
-    # - https://click.palletsprojects.com/en/8.1.x/api/#click.Context.forward
-    # - https://click.palletsprojects.com/en/8.1.x/api/#click.Context.invoke
 ```
+"Build" = [".spin/cmds.py:build_extend"]
+```
+
+The `extend_cmd` decorator also accepts a `doc` argument, for setting the new command's `--help` description.
+The function documentation ("This version of build...") is also appended.
 
 ### Advanced: override Meson CLI
 

example_pkg/.spin/cmds.py

@@ -2,6 +2,7 @@
 
 import click
 
+import spin
 from spin import util
 
 
@@ -33,3 +34,15 @@ def example(flag, test, default_kwd=None):
 
     click.secho("\nTool config is:", fg="yellow")
     print(json.dumps(config["tool.spin"], indent=2))
+
+
+@click.option("-e", "--extra", help="Extra test flag", type=int)
+@util.extend_command(spin.cmds.meson.build)
+def build_ext(*, parent_callback, extra=None, **kwargs):
+    """
+    This version of build also provides the EXTRA flag, that can be used
+    to specify an extra integer argument.
+    """
+    print(f"Preparing for build with {extra=}")
+    parent_callback(**kwargs)
+    print("Finalizing build...")

example_pkg/pyproject.toml

@@ -43,7 +43,7 @@ package = 'example_pkg'
   "spin.cmds.meson.gdb",
   "spin.cmds.meson.lldb"
 ]
-"Extensions" = [".spin/cmds.py:example"]
+"Extensions" = [".spin/cmds.py:example", ".spin/cmds.py:build_ext"]
 "Pip" = [
   "spin.cmds.pip.install"
 ]

spin/__main__.py

@@ -131,6 +131,8 @@ def group(ctx):
     }
     cmd_default_kwargs = toml_config.get("tool.spin.kwargs", {})
 
+    custom_module_cache = {}
+
     for section, cmds in config_cmds.items():
         for cmd in cmds:
             if cmd not in commands:
@@ -147,11 +149,17 @@ def group(ctx):
                 else:
                     try:
                         path, func = cmd.split(":")
-                        spec = importlib.util.spec_from_file_location(
-                            "custom_mod", path
-                        )
-                        mod = importlib.util.module_from_spec(spec)
-                        spec.loader.exec_module(mod)
+
+                        if path not in custom_module_cache:
+                            spec = importlib.util.spec_from_file_location(
+                                "custom_mod", path
+                            )
+                            mod = importlib.util.module_from_spec(spec)
+                            spec.loader.exec_module(mod)
+                            custom_module_cache[path] = mod
+                        else:
+                            mod = custom_module_cache[path]
+
                     except FileNotFoundError:
                         print(
                             f"!! Could not find file `{path}` to load custom command `{cmd}`.\n"

spin/cmds/util.py

@@ -2,10 +2,12 @@
     annotations,  # noqa: F401  # TODO: remove once only >3.8 is supported
 )
 
+import copy
 import os
 import shlex
 import subprocess
 import sys
+from collections.abc import Callable
 
 import click
 
@@ -98,3 +100,68 @@ def get_commands():
        ``commands`` key.
     """
     return click.get_current_context().meta["commands"]
+
+
+Decorator = Callable[[Callable], Callable]
+
+
+def extend_command(cmd: click.Command, doc: str | None = None) -> Decorator:
+    """This is a decorator factory.
+
+    The resulting decorator lets the user derive their own command from `cmd`.
+    The new command can support arguments not supported by `cmd`.
+
+    Parameters
+    ----------
+    cmd : click.Command
+        Command to extend.
+    doc : str
+        Replacement docstring.
+        The wrapped function's docstring is also appended.
+
+    Examples
+    --------
+
+    @click.option("-e", "--extra", help="Extra test flag")
+    @util.extend_cmd(
+        spin.cmds.meson.build
+    )
+    @extend_cmd(spin.cmds.meson.build)
+    def build(*args, constant=None, **kwargs):
+        '''
+        Some extra documentation related to the constant flag.
+        '''
+        ...
+        ctx.forward(spin.cmds.meson.build, *args, **kwargs)
+        ...
+
+    """
+    my_cmd = copy.copy(cmd)
+
+    # This is necessary to ensure that added options do not leak
+    # to the original command
+    my_cmd.params = copy.deepcopy(cmd.params)
+
+    def decorator(user_func: Callable) -> click.Command:
+        def callback_with_parent_callback(ctx, *args, **kwargs):
+            """Wrap the user callback to receive a
+            `parent_callback` keyword argument, containing the
+            callback from the originally wrapped command."""
+
+            def parent_cmd(*user_args, **user_kwargs):
+                ctx.invoke(cmd.callback, *user_args, **user_kwargs)
+
+            return user_func(*args, parent_callback=parent_cmd, **kwargs)
+
+        my_cmd.callback = click.pass_context(callback_with_parent_callback)
+
+        if doc is not None:
+            my_cmd.help = doc
+        my_cmd.help = (my_cmd.help or "") + "\n\n" + (user_func.__doc__ or "")
+        my_cmd.help = my_cmd.help.strip()
+
+        my_cmd.name = user_func.__name__.replace("_", "-")
+
+        return my_cmd
+
+    return decorator

spin/tests/test_extend_command.py

@@ -0,0 +1,74 @@
+import click
+import pytest
+
+from spin import cmds
+from spin.cmds.util import extend_command
+
+from .testutil import get_usage, spin
+
+
+def test_override_add_option():
+    @click.option("-e", "--extra", help="Extra test flag")
+    @extend_command(cmds.meson.build)
+    def build_ext(*, parent_callback, extra=None, **kwargs):
+        pass
+
+    assert "--extra" in get_usage(build_ext)
+    assert "--extra" not in get_usage(cmds.meson.build)
+
+
+def test_doc_setter():
+    @click.option("-e", "--extra", help="Extra test flag")
+    @extend_command(cmds.meson.build)
+    def build_ext(*, parent_callback, extra=None, **kwargs):
+        """
+        Additional docstring
+        """
+        pass
+
+    assert "Additional docstring" in get_usage(build_ext)
+    assert "Additional docstring" not in get_usage(cmds.meson.build)
+
+    @extend_command(cmds.meson.build, doc="Hello world")
+    def build_ext(*, parent_callback, extra=None, **kwargs):
+        """
+        Additional docstring
+        """
+        pass
+
+    doc = get_usage(build_ext)
+    assert "Hello world\n" in doc
+    assert "\n  Additional docstring" in doc
+
+
+def test_ext_additional_args():
+    @click.option("-e", "--extra", help="Extra test flag", type=int)
+    @extend_command(cmds.meson.build)
+    def build_ext(*, parent_callback, extra=None, **kwargs):
+        """
+        Additional docstring
+        """
+        assert extra == 5
+
+    ctx = build_ext.make_context(
+        None,
+        [
+            "--extra=5",
+        ],
+    )
+    ctx.forward(build_ext)
+
+    # And ensure that option didn't leak into original command
+    with pytest.raises(click.exceptions.NoSuchOption):
+        cmds.meson.build.make_context(
+            None,
+            [
+                "--extra=5",
+            ],
+        )
+
+
+def test_cli_additional_arg(example_pkg):
+    p = spin("build-ext", "--extra=3")
+    assert b"Preparing for build with extra=3" in p.stdout
+    assert b"meson compile" in p.stdout

spin/tests/testutil.py

@@ -38,3 +38,8 @@ def stdout(p):
 
 def stderr(p):
     return p.stderr.decode("utf-8").strip()
+
+
+def get_usage(cmd):
+    ctx = cmd.make_context(None, [])
+    return cmd.get_help(ctx)