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

【问题标题】：How to decorate a Python function with a context manager?如何用上下文管理器装饰 Python 函数？
【发布时间】：2021-08-24 07:11:31
【问题描述】：

鉴于现有的上下文管理器，我想要一个函数装饰器，它会导致在 with-block 内发生执行。

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

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

await f()

async def f():
    ...

async with contextmanager:
    await f()

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

像decorate 这样的便利实用程序是否已经存在于标准库的某个地方？

【问题讨论】：

标签： python python-decorators contextmanager

【解决方案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

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

【讨论】：