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

В контексте питона click Приложение CLI, я хотел бы запустить подкоманду внутри диспетчера контекста, которая будет настроена в команде более высокого уровня. Как это можно сделать с 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Не то чтобы я хотел решить эту проблему в частности. Я знаю, что могу повторить один и тот же код во всех командах и добиться того же эффекта, который я уже делал в моем реальном случае использования. Мой вопрос именно о том, что я могу сделать только в основной функции, которая будет запускать все прозрачные подкоманды в диспетчере контекста.

2 ответа

Новое решение

Наконец-то сделал это!!!

Вот способ пойти:

  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,

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,

Этот вариант использования изначально поддерживается в Click from v8.0 с помощью ctx.with_resource(context_manager)

https://click.palletsprojects.com/en/8.0.x/api/#click.Context.with_resource

В расширенной документации Click есть рабочий пример.

https://click.palletsprojects.com/en/8.0.x/advanced/#managing-resources

Другие вопросы по тегам