Q:

Run a subcommand inside a context manager

In the context of a python click CLI application, I would like to run a subcommand inside of a context manager that would be setup in a higher level command. How is it possible to do that with click? My pseudo-code looks something like:

import click

from contextlib import contextmanager

@contextmanager
def database_context(db_url):
    try:
        print(f'setup db connection: {db_url}')
        yield
    finally:
        print('teardown db connection')

@click.group
@click.option('--db',default='local')
def main(db):
    print(f'running command against {db} database')
    db_url = get_db_url(db)
    connection_manager = database_context(db_url)
    # here come the mysterious part that makes all subcommands
    # run inside the connection manager

@main.command
def do_this_thing()
    print('doing this thing')

@main.command
def do_that_thing()
    print('doing that thing')

And this would be called like:
> that_cli do_that_thing
running command against local database
setup db connection: db://user:pass@localdb:db_name
doing that thing
teardown db connection

> that_cli --db staging do_this_thing
running command against staging database
setup db connection: db://user:pass@123.456.123.789:db_name
doing this thing
teardown db connection

Edit: note that the above example is forged to better illustrate the missing functionality of click, not that I want to solve this problem in particular. I know I could repeat the same code in all commands and achieve the same effect, which I already do in my real use case. My question is precisely on what could I do only in the main function, that would run all transparently subcommands in a context manager.

A:

Decorating commands

Define a context manager decorator using contextlib.ContextDecorator
Use click.pass_context decorator on main(), so you can explore click context
Create an instance db_context of the context manager
Iterate on commands defined for group main using ctx.command.commands
For each command, replace the original callback (function called by the command) with the same callback decorated with the context manager db_context(cmd)

This way you will programmatically modify each command to behave just like:
@main.command()
@db_context
def do_this_thing():
    print('doing this thing')

But without requiring to change your code beyond your function main().
See the code below for a working example:
import click
from contextlib import ContextDecorator

class Database_context(ContextDecorator):
    """Decorator context manager."""

    def __init__(self, db_url):
        self.db_url = db_url

    def __enter__(self):
        print(f'setup db connection: {self.db_url}')

    def __exit__(self, type, value, traceback):
        print('teardown db connection')

@click.group() 
@click.option('--db', default='local')
@click.pass_context
def main(ctx, db):

    print(f'running command against {db} database')
    db_url = db  # get_db_url(db)

# here come the mysterious part that makes all subcommands
# run inside the connection manager

    db_context = Database_context(db_url)           # Init context manager decorator
    for name, cmd in ctx.command.commands.items():  # Iterate over main.commands
        cmd.allow_extra_args = True                 # Seems to be required, not sure why
        cmd.callback = db_context(cmd.callback)     # Decorate command callback with context manager

@main.command()
def do_this_thing():
    print('doing this thing')

@main.command()
def do_that_thing():
    print('doing that thing')

if __name__ == "__main__":
    main()

It does what you describe in your question, hope it will work as expected in real code.

Using click.pass_context
This code below will give you an idea of how to do it using click.pass_context.
import click
from contextlib import contextmanager

@contextmanager
def database_context(db_url):
    try:
        print(f'setup db connection: {db_url}')
        yield
    finally:
        print('teardown db connection')

@click.group()
@click.option('--db',default='local')
@click.pass_context
def main(ctx, db):
    ctx.ensure_object(dict)
    print(f'running command against {db} database')
    db_url = db #get_db_url(db)
    # Initiate context manager
    ctx.obj['context'] = database_context(db_url)

@main.command()
@click.pass_context
def do_this_thing(ctx):
    with ctx.obj['context']:
        print('doing this thing')

@main.command()
@click.pass_context
def do_that_thing(ctx):
    with ctx.obj['context']:
        print('doing that thing')

if __name__ == "__main__":
    main(obj={})

Another solution to avoid explicit with statement could be passing the context manager as a decorator using contextlib.ContextDecorator, but it would likely be more complex to setup with click.

