CODEWITHSUNDEEP
pythoncommand-line-interfacepython-click
Mike
Mike

Reputation: 872

Python click different named arguments based on number of arguments

How can I achieve the following synopsis using the Python click library?

Usage: app CMD [OPTIONS] [FOO] [BAR]
       app CMD [OPTIONS] [FOOBAR]

I can't figure out whether I am able to pass two different sets of named argument for the same command based on the number of given arguments. That is, if only one argument was passed it's foobar, but if two arguments were passed, they are foo and bar.

The code representation of such implementation would look something like this (provided you could use function overload, which you can't)

@click.command()
@click.argument('foo', required=False)
@click.argument('bar', required=False)
def cmd(foo, bar):
    # ...

@click.command()
@click.argument('foobar', required=False)
def cmd(foobar):
    # ...

Upvotes: 1

Views: 1141

Answers (1)

Stephen Rauch
Stephen Rauch

Reputation: 49794

You can add multiple command handlers with a different number of arguments for each by creating a custom click.Command class. There is some ambiguity around which of the command handlers would best be called if parameters are not strictly required, but that can be mostly dealt with by using the first signature that fits the command line passed.

Custom Class

class AlternateArgListCmd(click.Command):

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

    def alternate_arglist(self, *args, **kwargs):
        from click.decorators import command as cmd_decorator

        def decorator(f):
            command = cmd_decorator(*args, **kwargs)(f)
            self.alternate_arglist_handlers.append((command, command))

            # verify we have no options defined and then copy options from base command
            options = [o for o in command.params if isinstance(o, click.Option)]
            if options:
                raise click.ClickException(
                    f'Options not allowed on {type(self).__name__}: {[o.name for o in options]}')
            command.params.extend(o for o in self.params if isinstance(o, click.Option))
            return command

        return decorator

    def make_context(self, info_name, args, parent=None, **extra):
        """Attempt to build a context for each variant, use the first that succeeds"""
        orig_args = list(args)
        for handler, handler_super in self.alternate_arglist_handlers:
            args[:] = list(orig_args)
            self.alternate_self = handler
            try:
                return handler_super.make_context(info_name, args, parent, **extra)
            except click.UsageError:
                pass
            except:
                raise

        # if all alternates fail, return the error message for the first command defined
        args[:] = orig_args
        return super().make_context(info_name, args, parent, **extra)

    def invoke(self, ctx):
        """Use the callback for the appropriate variant"""
        if self.alternate_self.callback is not None:
            return ctx.invoke(self.alternate_self.callback, **ctx.params)
        return super().invoke(ctx)

    def format_usage(self, ctx, formatter):
        """Build a Usage for each variant"""
        prefix = "Usage: "
        for _, handler_super in self.alternate_arglist_handlers:
            pieces = handler_super.collect_usage_pieces(ctx)
            formatter.write_usage(ctx.command_path, " ".join(pieces), prefix=prefix)
            prefix = " " * len(prefix)

Using the Custom Class:

To use the custom class, pass it as the cls argument to the click.command decorator like:

@click.command(cls=AlternateArgListCmd)
@click.argument('foo')
@click.argument('bar')
def cli(foo, bar):
    ...

Then use the alternate_arglist() decorator on the command to add another command handler with different arguments.

@cli.alternate_arglist()
@click.argument('foobar')
def cli_one_param(foobar):
    ...

How does this work?

This works because click is a well designed OO framework. The @click.command() decorator usually instantiates a click.Command object but allows this behavior to be over ridden with the cls parameter. So it is a relatively easy matter to inherit from click.Command in our own class and over ride the desired methods.

In this case we add a new decorator method: alternate_arglist(), and override three methods: make_context(), invoke() & format_usage(). The overridden make_context() method checks to see which of the command handler variants matches the number of args passed, the overridden invoke() method is used to call the appropriate command handler variant and the overridden format_usage() is used to create the help message showing the various usages.

Test Code:

import click


@click.command(cls=AlternateArgListCmd)
@click.argument('foo')
@click.argument('bar')
@click.argument('baz')
@click.argument('bing', required=False)
@click.option('--an-option', default='empty')
def cli(foo, bar, baz, bing, an_option):
    """Best Command Ever!"""
    if bing is not None:
        click.echo(f'foo bar baz bing an-option: {foo} {bar} {baz} {bing} {an_option}')
    else:
        click.echo(f'foo bar baz an-option: {foo} {bar} {baz} {an_option}')


@cli.alternate_arglist()
@click.argument('foo')
@click.argument('bar')
def cli_two_param(foo, bar, an_option):
    click.echo(f'foo bar an-option: {foo} {bar} {an_option}')


@cli.alternate_arglist()
@click.argument('foobar', required=False)
def cli_one_param(foobar, an_option):
    click.echo(f'foobar an-option: {foobar} {an_option}')


if __name__ == "__main__":
    commands = (
        '',
        'p1',
        'p1 p2 --an-option=optional',
        'p1 p2 p3',
        'p1 p2 p3 p4 --an-option=optional',
        'p1 p2 p3 p4 p5',
        '--help',
    )

    import sys, time
    time.sleep(1)
    print('Click Version: {}'.format(click.__version__))
    print('Python Version: {}'.format(sys.version))
    for cmd in commands:
        try:
            time.sleep(0.1)
            print('-----------')
            print('> ' + cmd)
            time.sleep(0.1)
            cli(cmd.split())

        except BaseException as exc:
            if str(exc) != '0' and \
                    not isinstance(exc, (click.ClickException, SystemExit)):
                raise

Test Results:

Click Version: 7.1.2
Python Version: 3.8.5 (tags/v3.8.5:580fbb0, Jul 20 2020, 15:57:54) [MSC v.1924 64 bit (AMD64)]
-----------
>
foobar an-option: None empty
-----------
> p1
foobar an-option: p1 empty
-----------
> p1 p2 --an-option=optional
foo bar an-option: p1 p2 optional
-----------
> p1 p2 p3
foo bar baz an-option: p1 p2 p3 empty
-----------
> p1 p2 p3 p4 --an-option=optional
foo bar baz bing an-option: p1 p2 p3 p4 optional
-----------
> p1 p2 p3 p4 p5
Usage: test_code.py [OPTIONS] FOO BAR BAZ [BING]
       test_code.py [OPTIONS] FOO BAR
       test_code.py [OPTIONS] [FOOBAR]
Try 'test_code.py --help' for help.

Error: Got unexpected extra argument (p5)
-----------
> --help
Usage: test_code.py [OPTIONS] FOO BAR BAZ [BING]
       test_code.py [OPTIONS] FOO BAR
       test_code.py [OPTIONS] [FOOBAR]

  Best Command Ever!

Options:
  --an-option TEXT
  --help            Show this message and exit.

Upvotes: 1

Related Questions