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

Question

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.

I was hoping for something like (thanks Stephen Rauch for the initial idea):

def do_upgrade():
    print "Performing upgrade"

def bypass_upgrade_check(func):
    setattr(func, "do_upgrade_check", False)
    return func

@click.group()
@click.pass_context
def common(ctx):
    sub_cmd = ctx.command.commands[ctx.invoked_subcommand]
    if getattr(sub_cmd, "do_upgrade_check", True):
        do_upgrade()

@bypass_upgrade_check
@common.command()
def top_cmd1():
    # don't run do_upgrade() on top level command
    pass

@common.command()
def top_cmd2():
    # DO run do_upgrade() on top level command
    pass

@common.group()
def sub_cmd_group():
    pass

@bypass_upgrade_check
@sub_cmd.command()
def sub_cmd1():
    # don't run do_upgrade() on second-level command
    pass

@sub.command()
def sub_cmd2():
    # DO run do_upgrade() on second-level command
    pass

Unfortunately, this only works for the top-level commands, since ctx.invoked_subcommand refers to sub_cmd_group and not sub_cmd1 or sub_cmd2.

Is there a way to recursively search through sub-commands, or perhaps using a custom Group to be able to achieve this functionality with both top-level commands as well as subcommands?

Answer 1

One way to solve this is to build a custom decorator that pairs with a custom click.Group class:

Custom Decorator Builder:

def make_exclude_hook_group(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 invoke(self, ctx):
            """ group invoke which hooks context invoke """
            invoke = ctx.invoke

            def ctx_invoke(*args, **kwargs):
                """ monkey patched context invoke """
                sub_cmd = ctx.command.commands[ctx.invoked_subcommand]
                if not isinstance(sub_cmd, click.Group) and \
                        getattr(sub_cmd, hook_attr_name, True):
                    # invoke the callback
                    callback()
                return invoke(*args, **kwargs)

            ctx.invoke = ctx_invoke

            return super(HookGroup, self).invoke(ctx)

        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 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_check = make_exclude_hook_group(do_upgrade)

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

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

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

@bypass_upgrade_check
@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 monkey patch click.Context.invoke() of our context and if the command that is about to be executed has not been decorated, we call the callback.

Test Code:

import click

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

bypass_upgrade_check = make_exclude_hook_group(do_upgrade)

@click.group(cls=bypass_upgrade_check())
@click.pass_context
def cli(ctx):
    pass

@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')

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)]
-----------
> top_cmd1
cmd1
-----------
> top_cmd2
Performing upgrade
cmd2
-----------
> sub_cmd_group sub_cmd1
sub_cmd_group
sub_cmd1
-----------
> sub_cmd_group sub_cmd2
Performing upgrade
sub_cmd_group
sub_cmd2
-----------
> --help
Usage: test.py [OPTIONS] COMMAND [ARGS]...

Options:
  --help  Show this message and exit.

Commands:
  sub_cmd_group
  top_cmd1
  top_cmd2