Запустите подкоманду внутри контекстного менеджера

Question

В контексте приложения CLI для Python click я хотел бы запустить подкоманду внутри диспетчера контекста, которая будет настроена в команде более высокого уровня.Как это можно сделать с click?Мой псевдокод выглядит примерно так:

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

И это будет называться так:

> 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

Редактировать: обратите внимание, что приведенный выше пример подделан для лучшей иллюстрации отсутствующей функциональности click, не то чтобы я хотел решить эту проблему в частности.Я знаю, что могу повторить один и тот же код во всех командах и добиться того же эффекта, который я уже делал в моем реальном случае использования.Мой вопрос как раз о том, что я могу сделать только в основной функции, которая будет запускать все прозрачные подкоманды в менеджере контекста.

Answer 1

Команды декорирования

1. Определение менеджера контекста Декоратор с использованием contextlib.ContextDecorator
2. Использование click.pass_context декоратор в main(), так что вы можете исследовать контекст клика
3. Создать экземпляр db_context диспетчера контекста
4. Выполнять итерации по командам, определенным для группы main, используя ctx.command.commands
5. Для каждой команды замените исходный обратный вызов (функция, вызываемую командой) тем же обратным вызовом, украшенным менеджером контекста db_context(cmd)

Таким образом, вы программно измените каждую команду на поведениепросто как:

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

Но без необходимости изменять ваш код вне вашей функции main().

Рабочий код приведен ниже для рабочего примера:

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()

Он делает то, что вы описываете в своем вопросе, надеюсь, что он будет работать так, как ожидалось в реальном коде.

---

Использование click.pass_context

Этот код ниже даст вамидея, как это сделать, используя 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={})

Другое решениево избежание явного выражения with оператор может передавать диспетчер контекста в качестве декоратора, используя contextlib.ContextDecorator, но его, вероятно, будет сложнее настроить с помощью click.