Search code examples
pythonpython-click

Click: how do I apply an action to all commands and subcommands but allow a command to opt-out (part duex)?


Building on my initial question, I'd like to be able to have the body of a parent group run before I run my callback.

I have a case where I'd like to automatically run a common function, check_upgrade(), for most of my click commands and sub-commands, but there are a few cases where I don't want to run it. I was thinking I could have a decorator that one can add (e.g. @bypass_upgrade_check) for commands where check_upgrade() should not run.

For example:

def do_upgrade():
    print("Performing upgrade")

bypass_upgrade_check = make_exclude_hook_group(do_upgrade)

@click.group(cls=bypass_upgrade_check())
@click.option('--arg1', default=DFLT_ARG1)
@click.option('--arg2', default=DFLT_ARG2)
@click.pass_context
def cli(ctx, arg1, arg2):
    config.call_me_before_upgrade_check(arg1, arg2)

@bypass_upgrade_check
@cli.command()
def top_cmd1():
    click.echo('cmd1')

@cli.command()
def top_cmd2():
    click.echo('cmd2')

@cli.group()
def sub_cmd_group():
    click.echo('sub_cmd_group')

@bypass_upgrade_check
@sub_cmd_group.command()
def sub_cmd1():
    click.echo('sub_cmd1')

@sub_cmd_group.command()
def sub_cmd2():
    click.echo('sub_cmd2')

I'd like things to function like explained in the initial question, but instead of executing do_upgrade() before executing the body of cli(), I'd like it to call:

cli() --> do_upgrade() --> top_cmd1()

for example. Or for a nested command:

cli() --> sub_cmd_group() --> do_upgrade() --> sub_cmd1()

So I guess another way to phrase the question is: is it possible to have the functionality from the original question, but have the callback be called right before the subcommand itself runs instead of being called before any of the Group blocks run?

The reason I need this is because the arguments passed in to the top-level CLI command indicate the server address to check for an upgrade. I need this information to process do_upgrade(). I can't pass this information directly to do_upgrade() because this server information is also used elsewhere in the application. I can query it from do_upgrade() with something like config.get_server().


Solution

  • In a similar fashion to the original question, one way to solve this is to build a custom decorator that pairs with a custom click.Group class. The added complication is to hook the Command.invoke() instead of the Group.invoke() so that the callback will be invoked immediately preceding the Command.invoke() and thus will be invoked after any Group.invoke():

    Custom Decorator Builder:

    import click
    
    def make_exclude_hook_command(callback):
        """ for any command that is not decorated, call the callback """
    
        hook_attr_name = 'hook_' + callback.__name__
    
        class HookGroup(click.Group):
            """ group to hook context invoke to see if the callback is needed"""
    
            def group(self, *args, **kwargs):
                """ new group decorator to make sure sub groups are also hooked """
                if 'cls' not in kwargs:
                    kwargs['cls'] = type(self)
                return super(HookGroup, self).group(*args, **kwargs)
    
            def command(self, *args, **kwargs):
                """ new command decorator to monkey patch command invoke """
    
                cmd = super(HookGroup, self).command(*args, **kwargs)
    
                def hook_command_decorate(f):
                    # decorate the command
                    ret = cmd(f)
    
                    # grab the original command invoke
                    orig_invoke = ret.invoke
    
                    def invoke(ctx):
                        """call the call back right before command invoke"""
                        parent = ctx.parent
                        sub_cmd = parent and parent.command.commands[
                            parent.invoked_subcommand]
                        if not sub_cmd or \
                                not isinstance(sub_cmd, click.Group) and \
                                getattr(sub_cmd, hook_attr_name, True):
                            # invoke the callback
                            callback()
                        return orig_invoke(ctx)
    
                    # hook our command invoke to command and return cmd
                    ret.invoke = invoke
                    return ret
    
                # return hooked command decorator
                return hook_command_decorate
    
        def decorator(func=None):
            if func is None:
                # if called other than as decorator, return group class
                return HookGroup
    
            setattr(func, hook_attr_name, False)
    
        return decorator
    

    Using the decorator builder:

    To use the decorator we first need to build the decorator like:

    bypass_upgrade = make_exclude_hook_command(do_upgrade)
    

    Then we need to use it as a custom class to click.group() like:

    @click.group(cls=bypass_upgrade())
    ...
    

    And finally, we can decorate any commands or sub-commands to the group that need to not use the callback like:

    @bypass_upgrade
    @my_group.command()
    def my_click_command_without_upgrade():
         ...
    

    How does this work?

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

    In this case, we build a decorator that sets an attribute on any click function that does not need the callback called. Then in our custom group, we overide both the group() and the command() decorators so that we can we monkey patch invoke() on the command and if the command that is about to be executed has not been decorated, we call the callback.

    Test Code:

    def do_upgrade():
        click.echo("Performing upgrade")
    
    bypass_upgrade = make_exclude_hook_command(do_upgrade)
    
    @click.group(cls=bypass_upgrade())
    @click.pass_context
    def cli(ctx):
        click.echo('cli')
    
    @bypass_upgrade
    @cli.command()
    def top_cmd1():
        click.echo('cmd1')
    
    @cli.command()
    def top_cmd2():
        click.echo('cmd2')
    
    @cli.group()
    def sub_cmd_group():
        click.echo('sub_cmd_group')
    
    @bypass_upgrade
    @sub_cmd_group.command()
    def sub_cmd1():
        click.echo('sub_cmd1')
    
    @sub_cmd_group.command()
    def sub_cmd2():
        click.echo('sub_cmd2')
    
    
    if __name__ == "__main__":
        commands = (
            'top_cmd1',
            'top_cmd2',
            'sub_cmd_group sub_cmd1',
            'sub_cmd_group sub_cmd2',
            '--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)]
    -----------
    > sub_cmd_group sub_cmd2
    cli
    sub_cmd_group
    Performing upgrade
    sub_cmd2
    -----------
    > top_cmd1
    cli
    cmd1
    -----------
    > top_cmd2
    cli
    Performing upgrade
    cmd2
    -----------
    > sub_cmd_group sub_cmd1
    cli
    sub_cmd_group
    sub_cmd1
    -----------
    > sub_cmd_group sub_cmd2
    cli
    sub_cmd_group
    Performing upgrade
    sub_cmd2
    -----------
    > --help
    Usage: test.py [OPTIONS] COMMAND [ARGS]...
    
    Options:
      --arg1 TEXT
      --arg2 TEXT
      --help       Show this message and exit.
    
    Commands:
      sub_cmd_group
      top_cmd1
      top_cmd2