时间戳

首页 > 解决方案 > 如何用上下文管理器装饰 Python 函数？

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

问题描述

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

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

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

await f()

async def f():
    ...

async with contextmanager:
    await f()

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

decorate标准库中是否已经存在类似的便利实用程序？

标签： pythonpython-decoratorscontextmanager

解决方案

有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在相关位置

推荐阅读