args-overload-simpler

#!/bin/python3
from argparse import ArgumentParser, ArgumentError
from contextlib import contextmanager


@contextmanager
def scoped_setattr(obj, name, value):
    dne = object()

    try:
        original = getattr(obj, name)
    except AttributeError:
        original = dne

    setattr(obj, name, value)
    try:
        yield
    finally:
        if original is dne:
            delattr(obj, name)
        else:
            setattr(obj, name, original)


class OverloadParser(ArgumentParser):
    class _Parser(ArgumentParser):
        def exit(self, status = 0, message = None):
            raise ArgumentError(None, message)

        def _print_message(self, message, file = None):
            pass

        # XXX FIXME: do we need to handle --help here?

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._overloads = []

    def add_overload(self, **kwargs):
        # XXX FIXME: do we need to set add_help = False here?
        self._overloads += [_parser := self._Parser(**kwargs, parents = [self], add_help = False)]
        return _parser

    def _parse_known_args(self, args, namespace):
        errors = []

        for overload in self._overloads:
            try:
                parsed = overload.parse_args(args)
            except ArgumentError as error:
                errors += [(overload, str(error))]
                continue

            # found a match!
            for k, v in vars(parsed).items():
                setattr(namespace, k, v)
            return namespace, []

        # XXX FIXME: show actual usage patterns associated with each unknowns
        #parser.error("unrecognized arguments for every usage pattern: %s" % ", ".join("[%s]" % " ".join(x) for x in unknowns))
        return namespace, args


######


parser = ArgumentParser(prog='prog')
subparsers = parser.add_subparsers()

hello = subparsers.add_parser("hello")
hello.set_defaults(__command__ = "hello")
hello.add_argument("who", default="world", nargs="?")

with scoped_setattr(subparsers, "_parser_class", OverloadParser):
    overloads = subparsers.add_parser("overloads")
    overloads.set_defaults(__command__ = "overloads")

x = overloads.add_overload()
x.set_defaults(__which__ = "x")
x.add_argument("letter", choices=["a", "b", "c"])

y = overloads.add_overload()
y.set_defaults(__which__ = "y")
y.add_argument("number", type=int)

print(parser.parse_args())