Search code examples
pythoncommand-line-interfacedecoratorpython-decoratorspython-click

How do I return a parameter value from a decorator using Python's Click package?


In my CLI application, I have a configuration file containing information on how to connect to an API. This configuration can be set up for multiple environments: (test/prod/etc).

I want to ensure that the passed environment exists in the configuration file, so I've created a decorator and modeled it off of the confirmation_option.

This is a small example of the problem:

# -*- coding: utf-8 -*-

import sys
import click

TEST_DICT = {
    'production': 'ProdKey'
}


def environment_option(*param_decls, **attrs):
    """
    Check that the passed environment exists
    """
    def decorator(f):
        def callback(ctx, param, value):
            if not value:
                ctx.abort()
            try:
                TEST_DICT[value]
            except KeyError:
                click.secho("Bad environment provided: {}".format(value), fg='red')
                sys.exit(1)

        attrs.setdefault('callback', callback)
        return click.option(*(param_decls or ('--environment',)), **attrs)(f)
    return decorator


@click.group()
@click.pass_context
def cli(ctx):
    """My CLI"""

@cli.command()
@environment_option('-e', '--environment', help="Environment to associate this key with", required=True)
@click.pass_context
def show_keys(ctx, environment):
    """
    List the available keys for the selected environment
    """
    click.echo("Environment:", type(environment))


if __name__ == "__main__":
    sys.exit(cli())  # pragma: no cover

If I pass in an invalid environment, it works as expected.

$ python cli show-keys -e notreal
Bad environment provided: notreal

The problem I am having, is that a valid environment - one that doesn't trigger the error in the decorator - doesn't have a value when it gets back to the show_keys function. It is NoneType:

$ python cli show-keys -e production

Traceback (most recent call last):
  File "/home/devuser/.virtualenvs/mycli/bin/mycli", line 11, in <module>
    load_entry_point('mycli', 'console_scripts', 'mycli')()
  File "/home/devuser/.virtualenvs/mycli/lib/python3.6/site-packages/click/core.py", line 764, in __call__
    return self.main(*args, **kwargs)
  File "/home/devuser/.virtualenvs/mycli/lib/python3.6/site-packages/click/core.py", line 717, in main
    rv = self.invoke(ctx)
  File "/home/devuser/.virtualenvs/mycli/lib/python3.6/site-packages/click/core.py", line 1137, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/devuser/.virtualenvs/mycli/lib/python3.6/site-packages/click/core.py", line 1137, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/devuser/.virtualenvs/mycli/lib/python3.6/site-packages/click/core.py", line 956, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/devuser/.virtualenvs/mycli/lib/python3.6/site-packages/click/core.py", line 555, in invoke
    return callback(*args, **kwargs)
  File "/home/devuser/.virtualenvs/mycli/lib/python3.6/site-packages/click/decorators.py", line 17, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "/home/devuser/repositories/mycli/mycli/mycli/commands/config_cmds.py", line 60, in show_keys
    click.echo("Environment:", type(environment))
  File "/home/devuser/.virtualenvs/mycli/lib/python3.6/site-packages/click/utils.py", line 260, in echo
    file.write(message)
AttributeError: type object 'NoneType' has no attribute 'write'

How do I get the value of my environment back from this decorator instead of NoneType?


Solution

  • Python functions without a return statement, return None. So in the callback, you need to return the value of TEST_DICT[value]. Something like:

    def callback(ctx, param, value):
        if not value:
            ctx.abort()
        try:
            return TEST_DICT[value]
        except KeyError:
            ....
    

    Test Code:

    import sys
    import click
    
    TEST_DICT = {
        'production': 'ProdKey'
    }
    
    
    def environment_option(*param_decls, **attrs):
        """
        Check that the passed environment exists
        """
    
        def decorator(f):
            def callback(ctx, param, value):
                if not value:
                    ctx.abort()
                try:
                    return TEST_DICT[value]
                except KeyError:
                    click.secho("Bad environment provided: {}".format(value),
                                fg='red')
                    sys.exit(1)
    
            attrs.setdefault('callback', callback)
            return click.option(*(param_decls or ('--environment',)), **attrs)(f)
    
        return decorator
    
    
    @click.group()
    @click.pass_context
    def cli(ctx):
        """My CLI"""
    
    
    @cli.command('show-keys')
    @environment_option('-e', '--environment',
                        help="Environment to associate this key with",
                        required=True)
    @click.pass_context
    def show_keys(ctx, environment):
        """
        List the available keys for the selected environment
        """
        click.echo("Environment: {}".format(environment))
    
    
    if __name__ == "__main__":
        commands = (
            'show-keys -e production',
            'show-keys -e notreal',
            'show-keys --help',
            '--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
    

    Results:

    Click Version: 6.7
    Python Version: 3.6.3 (v3.6.3:2c5fed8, Oct  3 2017, 18:11:49) [MSC v.1900 64 bit (AMD64)]
    -----------
    > show-keys -e production
    Environment: ProdKey
    -----------
    > show-keys -e notreal
    Bad environment provided: notreal
    -----------
    > show-keys --help
    Usage: test.py show-keys [OPTIONS]
    
      List the available keys for the selected environment
    
    Options:
      -e, --environment TEXT  Environment to associate this key with  [required]
      --help                  Show this message and exit.
    -----------
    > --help
    Usage: test.py [OPTIONS] COMMAND [ARGS]...
    
      My CLI
    
    Options:
      --help  Show this message and exit.
    
    Commands:
      show-keys  List the available keys for the selected...
    -----------
    > 
    Usage: test.py [OPTIONS] COMMAND [ARGS]...
    
      My CLI
    
    Options:
      --help  Show this message and exit.
    
    Commands:
      show-keys  List the available keys for the selected...