cli: shell completions for bash and zsh

README.md

@@ -60,6 +60,19 @@ by zero-based position (use this for steps with no `name:`).
 | `--act-arg ARG` | extra argument passed through to `act` (repeatable) |
 | `-v`, `--verbose` | print the injection/act commands being run |
 
+### Shell completions
+
+`actbreak --completions bash` (or `zsh`) prints a completion script built
+from the argparse parser, so new flags show up without touching it:
+
+```bash
+# bash
+source <(actbreak --completions bash)
+
+# zsh
+source <(actbreak --completions zsh)
+```
+
 ### Examples
 
 ```

actbreak/cli.py

@@ -32,6 +32,17 @@
 """
 
 
+class PrintCompletionsAction(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        from .completions import generate_bash, generate_zsh
+
+        if values == "bash":
+            print(generate_bash(parser), end="")
+        elif values == "zsh":
+            print(generate_zsh(parser), end="")
+        parser.exit(0)
+
+
 def build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
         prog=PROG,
@@ -40,6 +51,12 @@ def build_parser() -> argparse.ArgumentParser:
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
     parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+    parser.add_argument(
+        "--completions",
+        choices=["bash", "zsh"],
+        action=PrintCompletionsAction,
+        help="print shell completion script",
+    )
 
     sub = parser.add_subparsers(dest="command", required=True)
 

actbreak/completions.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import argparse
+
+
+def _extract(parser: argparse.ArgumentParser) -> tuple[list[str], list[str], dict[str, list[str]]]:
+    top_flags: list[str] = []
+    commands: list[str] = []
+    cmd_flags: dict[str, list[str]] = {}
+
+    for action in parser._actions:
+        if isinstance(action, argparse._SubParsersAction):
+            for cmd, subparser in action.choices.items():
+                commands.append(cmd)
+                flags: list[str] = []
+                for subaction in subparser._actions:
+                    if subaction.option_strings and subaction.option_strings[0] != "-h":
+                        flags.extend(subaction.option_strings)
+                cmd_flags[cmd] = flags
+        elif action.option_strings and action.option_strings[0] != "-h":
+            top_flags.extend(action.option_strings)
+
+    return top_flags, commands, cmd_flags
+
+
+def generate_bash(parser: argparse.ArgumentParser) -> str:
+    top_flags, commands, cmd_flags = _extract(parser)
+
+    cases = []
+    for cmd, flags in cmd_flags.items():
+        if flags:
+            flags_str = " ".join(flags)
+            cases.append(
+                f'''        {cmd})
+            COMPREPLY=( $(compgen -W "{flags_str}" -- "$cur") )
+            return 0
+            ;;'''
+            )
+        else:
+            cases.append(
+                f'''        {cmd})
+            return 0
+            ;;'''
+            )
+
+    cases_str = "\n".join(cases)
+    commands_str = " ".join(commands)
+    top_flags_str = " ".join(top_flags)
+
+    return f'''\
+_actbreak() {{
+    local cur prev words cword
+    COMPREPLY=()
+    cur="${{COMP_WORDS[COMP_CWORD]}}"
+
+    local commands="{commands_str}"
+    local top_flags="{top_flags_str}"
+
+    if [[ ${{COMP_CWORD}} -eq 1 ]]; then
+        if [[ "$cur" == -* ]]; then
+            COMPREPLY=( $(compgen -W "$top_flags" -- "$cur") )
+        else
+            COMPREPLY=( $(compgen -W "$commands" -- "$cur") )
+        fi
+        return 0
+    fi
+
+    local command="${{COMP_WORDS[1]}}"
+    if [[ "$cur" == -* ]]; then
+        case "$command" in
+{cases_str}
+        esac
+    fi
+    return 0
+}}
+complete -F _actbreak actbreak
+'''
+
+
+def generate_zsh(parser: argparse.ArgumentParser) -> str:
+    top_flags, commands, cmd_flags = _extract(parser)
+
+    cases = []
+    for cmd, flags in cmd_flags.items():
+        if flags:
+            flag_args = " ".join(f'"{f}"' for f in flags)
+            cases.append(
+                f'''        {cmd})
+            _arguments {flag_args}
+            ;;'''
+            )
+        else:
+            cases.append(
+                f'''        {cmd})
+            ;;'''
+            )
+
+    cases_str = "\n".join(cases)
+    top_flag_args = " ".join(f'"{f}"' for f in top_flags)
+    commands_args = " ".join(f'"{c}"' for c in commands)
+
+    return f'''\
+#compdef actbreak
+
+_actbreak() {{
+    local context state state_descr line
+    typeset -A opt_args
+
+    _arguments -C \\
+        {top_flag_args} \\
+        '1: :->cmds' \\
+        '*::arg:->args'
+
+    case $state in
+        cmds)
+            _values "actbreak command" {commands_args}
+            ;;
+        args)
+            case $line[1] in
+{cases_str}
+            esac
+            ;;
+    esac
+}}
+
+_actbreak "$@"
+'''

tests/test_cli.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+import contextlib
+import io
 import unittest
 from unittest import mock
 
@@ -70,6 +72,33 @@ def test_version_flag(self):
             parser.parse_args(["--version"])
         self.assertEqual(ctx.exception.code, 0)
 
+    def _completions(self, shell):
+        parser = build_parser()
+        buf = io.StringIO()
+        with contextlib.redirect_stdout(buf):
+            with self.assertRaises(SystemExit) as ctx:
+                parser.parse_args(["--completions", shell])
+        self.assertEqual(ctx.exception.code, 0)
+        return buf.getvalue()
+
+    def test_completions_bash_covers_parser(self):
+        out = self._completions("bash")
+        self.assertIn("_actbreak() {", out)
+        for token in ("run", "resume", "clean", "--version", "--completions",
+                      "--break-before", "--break-after", "--break-on-failure",
+                      "--job", "--runtime", "--no-attach", "--act-arg",
+                      "-v", "--verbose"):
+            self.assertIn(token, out)
+
+    def test_completions_zsh_covers_parser(self):
+        out = self._completions("zsh")
+        self.assertIn("#compdef actbreak", out)
+        for token in ('"run"', '"resume"', '"clean"', '"--version"',
+                      '"--completions"', '"--break-before"', '"--break-after"',
+                      '"--break-on-failure"', '"--job"', '"--runtime"',
+                      '"--no-attach"', '"--act-arg"', '"-v"', '"--verbose"'):
+            self.assertIn(token, out)
+
     def test_missing_command_is_an_error(self):
         parser = build_parser()
         with self.assertRaises(SystemExit):