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

Question

给定一个现有的上下文管理器，我想要一个函数装饰器，使其在with块中执行。

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

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

await f()

async def f():
    ...

async with contextmanager:
    await f()

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

标准库中是否已经存在像decorate这样的便捷工具？

- benjimin

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

- Ron Serruya