command.py

from argparse import ArgumentParser

__all__ = ['execute', 'metacommand', 'command', 'argument']

_reserved_args = 'func', 'command'

_default_registry = {}


def execute(args, prog=None, commands=None):
    if commands is None:
        commands = _default_registry
    parsed = parse(commands, args, prog)
    kwargs = {k: v for k, v in vars(parsed).items() if k not in _reserved_args}
    parsed.func(**kwargs)


def parse(commands, args, prog):
    cli = ArgumentParser(prog=prog)
    root = cli.add_subparsers(dest='command')  # using reserved argument
    for command_name in commands:
        func, arguments, description = commands[command_name]
        parser = root.add_parser(command_name, help=description)
        if arguments is not None:
            for name_or_flags, options in arguments:
                if 'metavar' not in options and len(name_or_flags) == 1 and name_or_flags[0][0] != '-':
                    options = dict(metavar=name_or_flags[0], **options)
                    name_or_flags = [name_or_flags[0].replace('-', '_')]
                parser.add_argument(*name_or_flags, **options)

        parser.set_defaults(func=func)  # using reserved argument

    args = cli.parse_args(args)
    if args.command is None:
        cli.print_help()
        cli.exit(1)

    return args


def metacommand(registry):
    def command(arguments_or_callable=None, name=None):
        if callable(arguments_or_callable) and name is None:
            return command()(arguments_or_callable)

        arguments = arguments_or_callable

        def decorate(func):
            command_name = name
            if command_name is None:
                command_name = func.__name__.replace('_', '-')

            description = None
            if func.__doc__ is not None:
                description = func.__doc__.lstrip().splitlines()[0]

            if command_name in registry:
                raise ValueError('Duplicate command name')

            registry[command_name] = func, arguments, description

        return decorate
    return command


command = metacommand(_default_registry)


def argument(*name_or_flags, **options):
    if any(k in _reserved_args for k in name_or_flags):
        raise KeyError('Reserved argument name')
    return name_or_flags, options