"""
Shell auto-completion script generators for the CLI.
Introspects an argparse parser to enumerate all commands, subcommands, and
flags, then emits completion scripts for bash, zsh, and fish.
"""
import argparse
import re
from typing import Dict, List, Tuple
def _safe_identifier(prog: str) -> str:
"""Convert a prog name to a safe shell identifier (letters, digits, _)."""
return re.sub(r"[^a-zA-Z0-9]", "_", prog)
def _collect_commands(
parser: argparse.ArgumentParser,
) -> Dict[Tuple[str, ...], List[str]]:
"""Walk the parser tree and return {command_path: [flags]} mapping."""
result: Dict[Tuple[str, ...], List[str]] = {}
_walk_parser(parser, (), result)
return result
def _get_flags(parser: argparse.ArgumentParser) -> List[str]:
"""Extract all optional flags from a parser (excluding help)."""
flags = []
for action in parser._actions:
if isinstance(action, argparse._HelpAction):
continue
if isinstance(action, argparse._SubParsersAction):
continue
for opt in action.option_strings:
flags.append(opt)
return sorted(flags)
def _walk_parser(
parser: argparse.ArgumentParser,
path: Tuple[str, ...],
result: Dict[Tuple[str, ...], List[str]],
):
"""Recursively walk subparsers and collect command paths + flags."""
flags = _get_flags(parser)
result[path] = flags
for action in parser._actions:
if isinstance(action, argparse._SubParsersAction):
for name, subparser in action.choices.items():
_walk_parser(subparser, path + (name,), result)
def _get_subcommands(
commands: Dict[Tuple[str, ...], List[str]],
prefix: Tuple[str, ...],
) -> List[str]:
"""Get immediate subcommands of a given prefix."""
subs = set()
for path in commands:
if len(path) == len(prefix) + 1 and path[: len(prefix)] == prefix:
subs.add(path[-1])
return sorted(subs)
def generate_bash(
parser: argparse.ArgumentParser, prog: str = "imunify360-agent"
) -> str:
"""Generate a bash completion script."""
commands = _collect_commands(parser)
lines = []
lines.append(f"# bash completion for {prog}")
lines.append(f"# Auto-generated by {prog} completions bash")
lines.append("")
lines.append(f"_{_safe_identifier(prog)}_completions() {{")
lines.append(" local cur prev words cword")
lines.append(" if type _init_completion &>/dev/null; then")
lines.append(" _init_completion || return")
lines.append(" else")
lines.append(" COMPREPLY=()")
lines.append(' cur="${COMP_WORDS[COMP_CWORD]}"')
lines.append(' prev="${COMP_WORDS[COMP_CWORD-1]}"')
lines.append(' words=("${COMP_WORDS[@]}")')
lines.append(" cword=$COMP_CWORD")
lines.append(" fi")
lines.append("")
lines.append(" # Build the command path from words")
lines.append(' local cmd_path=""')
lines.append(" local i")
lines.append(" for (( i=1; i < cword; i++ )); do")
lines.append(' case "${words[i]}" in')
lines.append(" -*) continue ;;")
lines.append(
' *) cmd_path="${cmd_path:+${cmd_path} }${words[i]}" ;;'
)
lines.append(" esac")
lines.append(" done")
lines.append("")
lines.append(' case "$cmd_path" in')
# Sort by depth (deepest first) so more specific paths match first
all_paths = sorted(commands.keys(), key=lambda p: (-len(p), p))
for path in all_paths:
if not path:
continue
subs = _get_subcommands(commands, path)
flags = commands[path]
completions = " ".join(subs + flags)
pattern = " ".join(path)
lines.append(f' "{pattern}")')
lines.append(
f' COMPREPLY=($(compgen -W "{completions}" -- "$cur"))'
)
lines.append(" return ;;")
# Root level
root_subs = _get_subcommands(commands, ())
root_flags = commands.get((), [])
root_completions = " ".join(root_subs + root_flags)
lines.append(' "")')
lines.append(
f' COMPREPLY=($(compgen -W "{root_completions}" -- "$cur"))'
)
lines.append(" return ;;")
lines.append(" esac")
lines.append("}")
lines.append("")
lines.append(f"complete -F _{_safe_identifier(prog)}_completions {prog}")
lines.append("")
return "\n".join(lines)
def generate_zsh(
parser: argparse.ArgumentParser, prog: str = "imunify360-agent"
) -> str:
"""Generate a zsh completion script."""
commands = _collect_commands(parser)
func_name = f"_{_safe_identifier(prog)}"
lines = []
lines.append(f"#compdef {prog}")
lines.append(f"# zsh completion for {prog}")
lines.append(f"# Auto-generated by {prog} completions zsh")
lines.append("")
lines.append(f"{func_name}() {{")
lines.append(" local -a commands flags")
lines.append(" local cmd_path")
lines.append("")
lines.append(" # Build command path from words")
lines.append(" cmd_path=()")
lines.append(" for word in ${words[2,-1]}; do")
lines.append(" [[ $word == -* ]] && continue")
lines.append(' [[ $word == "$words[$CURRENT]" ]] && continue')
lines.append(" cmd_path+=($word)")
lines.append(" done")
lines.append("")
lines.append(' case "${cmd_path[*]}" in')
all_paths = sorted(commands.keys(), key=lambda p: (-len(p), p))
for path in all_paths:
if not path:
continue
subs = _get_subcommands(commands, path)
flags = commands[path]
pattern = " ".join(path)
lines.append(f' "{pattern}")')
if subs:
desc_list = " ".join(f'"{s}"' for s in subs)
lines.append(f" commands=({desc_list})")
if flags:
flag_list = " ".join(f'"{f}"' for f in flags)
lines.append(f" flags=({flag_list})")
lines.append(
" _describe 'command' commands -- flags && return"
if subs
else f" compadd -- {' '.join(flags)} && return"
)
lines.append(" ;;")
# Root level
root_subs = _get_subcommands(commands, ())
root_flags = commands.get((), [])
root_desc = " ".join(f'"{s}"' for s in root_subs)
lines.append(' "")')
lines.append(f" commands=({root_desc})")
if root_flags:
flag_list = " ".join(f'"{f}"' for f in root_flags)
lines.append(f" flags=({flag_list})")
lines.append(" _describe 'command' commands -- flags && return")
lines.append(" ;;")
lines.append(" esac")
lines.append("}")
lines.append("")
lines.append(f"{func_name}")
lines.append("")
return "\n".join(lines)
def generate_fish(
parser: argparse.ArgumentParser, prog: str = "imunify360-agent"
) -> str:
"""Generate a fish completion script."""
commands = _collect_commands(parser)
lines = []
lines.append(f"# fish completion for {prog}")
lines.append(f"# Auto-generated by {prog} completions fish")
lines.append("")
# For each command path, emit completions
# Fish uses conditions based on what subcommands have been entered
for path in sorted(commands.keys(), key=lambda p: (len(p), p)):
subs = _get_subcommands(commands, path)
flags = commands[path]
if not path:
# Root level subcommands
condition = (
"not __fish_seen_subcommand_from"
f" {' '.join(_get_subcommands(commands, ()))}"
)
for sub in subs:
lines.append(
f"complete -c {prog} -n '{condition}' -f -a '{sub}'"
)
for flag in flags:
if flag.startswith("--"):
lines.append(
f"complete -c {prog} -n '{condition}' -l '{flag[2:]}'"
)
elif flag.startswith("-"):
lines.append(
f"complete -c {prog} -n '{condition}' -s '{flag[1:]}'"
)
else:
# Build condition: must have seen parent commands but not children
seen_parts = []
for p in path:
seen_parts.append(f"__fish_seen_subcommand_from {p}")
condition = " && ".join(seen_parts)
child_subs = subs
if child_subs:
condition += (
" && not __fish_seen_subcommand_from"
f" {' '.join(child_subs)}"
)
for sub in subs:
lines.append(
f"complete -c {prog} -n '{condition}' -f -a '{sub}'"
)
for flag in flags:
if flag.startswith("--"):
lines.append(
f"complete -c {prog} -n '{condition}' -l '{flag[2:]}'"
)
elif flag.startswith("-"):
lines.append(
f"complete -c {prog} -n '{condition}' -s '{flag[1:]}'"
)
lines.append("")
return "\n".join(lines)
GENERATORS = {
"bash": generate_bash,
"zsh": generate_zsh,
"fish": generate_fish,
}
SUPPORTED_SHELLS = sorted(GENERATORS.keys())
def generate_completions(
parser: argparse.ArgumentParser,
shell: str,
prog: str = "imunify360-agent",
) -> str:
"""Generate completion script for the given shell.
Raises ValueError if shell is not supported.
"""
generator = GENERATORS.get(shell)
if generator is None:
raise ValueError(
f"Unsupported shell: {shell}. "
f"Supported shells: {', '.join(SUPPORTED_SHELLS)}"
)
return generator(parser, prog)
|