如何使用上下文管理器装饰 Python function？

Question

给定一个现有的上下文管理器，我想要一个 function 装饰器，它会导致在with -block 内执行。

这应该使以下两个代码块等效：

@decorate(contextmanager)
async def f():
   ...

await f()

async def f():
    ...

async with contextmanager:
    await f()

这对于在例如asyncio.Semaphore(n)上下文中包装函数很方便。

标准库中是否已经存在像decorate这样的便利实用程序？

Answer 1

有ContextDecorator ，您可以创建从它继承的上下文管理器并像这样使用它

from functools import wraps
from contextlib import ContextDecorator

class MyContext(ContextDecorator):
    def __enter__(self):
        print('in enter')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('in exit')

some_context = MyContext()

@some_context
def foo(x, y):
    print(x + y)

foo(5, 3)

但是，我不确定您是否可以将它用于您未编写的现有内容，或者它是否适用于异步内容，在这种情况下，您可以像这样编写一个装饰器：

from functools import wraps

# Some context manager
class MyContext:
    def __enter__(self):
        print('in enter')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('in exit')

# Decotator that gets an existing context manager and uses it
def context_decorator(context_manager):

    def inner(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            with context_manager:
                return func(*args, **kwargs)

        return wrapper

    return inner

some_context = MyContext()

@context_decorator(some_context)
def foo(x, y):
    print(x + y)

foo(5, 3)

in enter
8
in exit

附言。 请注意，对于 asyncio.Semephore 之类的内容，您需要修改代码以在相关位置使用async with和await