Generated Zsh/Bash Autocompletion for argparse.ArgumentParser

argparse_completion.py

import argparse
from itertools import chain


def get_parser_info(parser):
    opts = {}
    commands = {}
    if parser._positionals and parser._positionals._actions:
        for a in parser._positionals._actions:
            if isinstance(a, argparse._StoreAction) and not a.option_strings:
                if a.choices:
                    def apply_at_leaf(_commands, choices):
                        if _commands:
                            for k in _commands.keys():
                                apply_at_leaf(_commands[k][1], choices)
                        else:
                            for c in choices:
                                _commands[c] = (opts, {})
                    apply_at_leaf(commands, a.choices)
    if parser._optionals and parser._optionals._group_actions:
        for a in parser._optionals._group_actions:
            if a.help == argparse.SUPPRESS:
                continue
            if isinstance(a, argparse._HelpAction):
                continue
            if isinstance(a, argparse._StoreTrueAction):
                for o in a.option_strings:
                    if len(a.option_strings) < 2 or o.startswith('--'):
                        opts[o] = None
            elif isinstance(a, (argparse._StoreAction, argparse._AppendAction)):
                for o in a.option_strings:
                    if len(a.option_strings) < 2 or o.startswith('--'):
                        opts[o] = a.choices or []
    if not commands and parser._subparsers:
        for a in parser._subparsers._actions:
            if isinstance(a, argparse._SubParsersAction):
                for k, v in a.choices.items():
                    commands[k] = get_parser_info(v)
    return opts, commands


def get_subcommand_parser_info(parser, argv):
    opts, commands = get_parser_info(parser)
    while argv:
        while argv and argv[0] not in commands:
            argv.pop(0)
        if argv:
            opts, commands = commands[argv.pop(0)]
    return opts, commands


def get_valid(parser, argv):
    opts, commands = get_subcommand_parser_info(parser, list(argv))
    if argv and argv[-1].startswith('-'):
        o = argv[-1]
        if opts.get(o, None):
            return opts[o]
    return chain(opts.keys(), commands.keys())
   

# Example program called 'foo'
if __name__ == '__main__':
  import sys
  parser = argparse.ArgumentParser()
  parser.add_argument('--bar')
  subparsers = parser.add_subparsers()
  foo_parser = subparsers.add_parser('foo')
  foo_parser.add_argument('baz', choices=['a', 'b'])
  foo_parser.add_argument('foo2', choices=['c', 'd'])
  foo_parser.add_argument('--bar2')
  if (
      '-h' not in sys.argv and
      '--help' not in sys.argv and
      sys.argv[-1] == 'commands'
  ):
      for v in get_valid(parser, sys.argv[1:-1]):
          print(v)
      sys.exit(0)
  opts = parser.parse_args()

completion.sh

# source completion.sh

_foo()
{
    local cur prev opts base
    COMPREPLY=()
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"

    case "${prev}" in
        -h|--help) COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) ) ;;
        -*|--*)
            opts="$(${COMP_WORDS[*]}commands)"
            COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
        ;;
        *)
            opts="$(${COMP_WORDS[*]} commands)"
            COMPREPLY=( $(compgen -W "--help ${opts}" -- ${cur}) )
        ;;
    esac
}
complete -F _foo foo