Click: Run subcommand in context manager

[antoine-gallix]

If would like to run a subcommand inside of a context manager that would be setup in a higher 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:[email protected]:db_name
doing this thing
teardown db connection

[stefreak]

Interface could be:

@click.group
@click.option('--db',default='local')
@click.pass_context
def main(ctx, db):
    print(f'running command against {db} database')
    db_url = get_db_url(db)
    with database_context(db_url) as db_context:
        ctx.obj = db_context
        yield

@main.command
@pass_obj
def do_this_thing(db)
    print('doing this thing with db', db)

click could detect the usage of yield with inspect.isgeneratorfunction

[sammko]

I'm also looking for this. An ugly workaround is to run the click main function from within an ExitStack and somehow pass it inside (using obj for examle). You can then use enter_context inside the main function to enter the database context and then overwrite the ctx.obj.

from contextlib import ExitStack, contextmanager

class DatabaseConnection:
    pass

@contextmanager
def database_context(where):
    print(f"Connecting to db {where}")
    try:
        yield DatabaseConnection()
    finally:
        print(f"Cleaning up db connection {where}")

pass_db = click.make_pass_decorator(DatabaseConnection)

@click.group()
@click.pass_context
def main(ctx):
    estack = ctx.obj

    db: DatabaseConnection = estack.enter_context(database_context("url"))
    ctx.obj = db

@main.command()
@pass_db
def do_this(db):
    # Yay, i have a db
    print(db)


if __name__ == '__main__':
    with ExitStack() as es:
        main(obj=es)

[davidism]

This has already been possible with Context.call_on_close() as long as the resource provides a close() method (or just pass its __exit__() method. Instead of using a with block, create the resource, register its close() method, and add it to ctx.obj.

@click.group()
@click.option("--username")
def cli(username):
    ctx = click.get_current_context()
    ctx.ensure_object(dict)["db"] = db = connect_db(username)
    ctx.call_on_close(db.close)

1191 switches to using an ExitStack internally for both call_on_close() and a new with_resource() method that will return a resource after entering it into the stack. In 8.0 it will be possible to do this:

@click.group()
@click.option("--username")
def cli(username):
    ctx = click.get_current_context()
    ctx.ensure_object(dict)["db"] = ctx.with_resource(connect_db(username))