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

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

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

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

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

await f()

async def f():
    ...

async with contextmanager:
    await f()

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

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

解决方法

有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_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,3)

in enter
8
in exit

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