运行 上下文管理器中的子命令

Question

在 python 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 缺少的功能而伪造的，并不是我要特别解决这个问题。我知道我可以在所有命令中重复相同的代码并达到相同的效果，我已经在实际用例中这样做了。我的问题恰恰是我可以在主函数中做什么 only，这将 运行 上下文管理器中的所有透明子命令。

Answer 1

修饰命令

1. 使用contextlib.ContextDecorator
定义上下文管理器装饰器
2. 在 main() 上使用 click.pass_context 装饰器，这样您就可以探索点击上下文
3. 创建上下文管理器的实例db_context
4. 使用 ctx.command.commands
迭代为组 main 定义的命令
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 进行设置可能会更复杂。

Answer 2

此用例在 v8.0 的 Click 中得到原生支持，方法是使用 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