在 Python 中充当装饰器和上下文管理器的函数？

Posted 2023-02-23

【中文标题】在 Python 中充当装饰器和上下文管理器的函数？

【英文标题】：Function acting as both decorator and context manager in Python?

【发布时间】：2012-03-02 02:02:48

【问题描述】：

这可能有点太过分了，但主要是出于好奇..

是否有可能有一个可调用的对象（函数/类）同时充当 上下文管理器和装饰器：

def xxx(*args, **kw):
    # or as a class

@xxx(foo, bar)
def im_decorated(a, b):
    print('do the stuff')

with xxx(foo, bar):
    print('do the stuff')

【问题讨论】：

你能举例说明它会做什么吗？这两个功能有区别吗？

这基本上是在测试套件中设置夹具的两种替代方式。

【参考方案1】：

从 Python 3.2 开始，对此的支持甚至包含在标准库中。从类contextlib.ContextDecorator 派生可以很容易地编写可以用作装饰器或上下文管理器的类。这个功能可以很容易地向后移植到 Python 2.x——这是一个基本的实现：

class ContextDecorator(object):
    def __call__(self, f):
        @functools.wraps(f)
        def decorated(*args, **kwds):
            with self:
                return f(*args, **kwds)
        return decorated

从此类派生上下文管理器，并像往常一样定义 __enter__() 和 __exit__() 方法。

【讨论】：

如果必须使用Python2可以使用contextlib2：contextlib2.readthedocs.org/en/latest

这可能会有所帮助：coderwall.com/p/0lk6jg/…

使用 python 3.8.5，添加或删除行 @functools.wraps(f) 不会改变任何内容，上下文管理器和无参数装饰器都可以工作。这是为什么呢？

@Guimute 请阅读functools.wraps() 的文档以了解它的作用。它不会改变包装函数工作的方式，但会改变它在检查时看起来的方式。例如，包装器继承了文档字符串和被包装函数的名称。

【参考方案2】：

这是一个例子：

class ContextDecorator(object):
    def __init__(self, foo, bar):
        self.foo = foo
        self.bar = bar
        print("init", foo, bar)

    def __call__(self, f):
        print("call")
        def wrapped_f():
            print("about to call")
            f()
            print("done calling")
        return wrapped_f

    def __enter__(self):
        print("enter")

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

with ContextDecorator(1, 2):
    print("with")

@ContextDecorator(3, 4)
def sample():
    print("sample")

sample()

打印出来：

init 1 2
enter
with
exit
init 3 4
call
about to call
sample
done calling

【讨论】：

当用作装饰器时，它的作用与上下文管理器不同，这似乎是 OP 的意图。 （请参阅 OP 的评论——“这基本上是在测试套件中设置固定装置的两种替代方法。”）

当然，重点是展示如何让一个班级同时完成这两项工作。我会把它留给 OP 来定制它以实现预期的功能。

【参考方案3】：

class Decontext(object):
    """
    makes a context manager also act as decorator
    """
    def __init__(self, context_manager):
        self._cm = context_manager
    def __enter__(self):
        return self._cm.__enter__()
    def __exit__(self, *args, **kwds):
        return self._cm.__exit__(*args, **kwds)
    def __call__(self, func):
        def wrapper(*args, **kwds):
            with self:
                return func(*args, **kwds)
        return wrapper

现在你可以这样做了：

mydeco = Decontext(some_context_manager)

这两者都允许

@mydeco
def foo(...):
    do_bar()

foo(...)

和

with mydeco:
    do_bar()

【讨论】：

如果上下文管理器接受参数，你将如何将参数传递给装饰器？

@Warz 上下文管理器已在 Decontext(some_context_manager) 期间创建。 __enter__ 和 __exit__ 方法是预定义的。

使用带有参数的Decontext 的相关问题：Passing arguments to decontext decorator

【参考方案4】：

在 Python 3.2+ 中，您可以使用 @contextlib.contextmanager 定义一个也是装饰器的上下文管理器。

来自文档：

contextmanager() 使用ContextDecorator，因此它创建的上下文管理器可以用作装饰器以及with 语句中

示例用法：

>>> from contextlib import contextmanager
>>> @contextmanager
... def example_manager(message):
...     print('Starting', message)
...     try:
...         yield
...     finally:
...         print('Done', message)
... 
>>> with example_manager('printing Hello World'):
...     print('Hello, World!')
... 
Starting printing Hello World
Hello, World!
Done printing Hello World
>>> 
>>> @example_manager('running my function')
... def some_function():
...     print('Inside my function')
... 
>>> some_function()
Starting running my function
Inside my function
Done running my function

【讨论】：

如果example_manager 产生了结果，当用作装饰器时，我们如何访问它？

@coler-j 你得到答案了吗？

【参考方案5】：

虽然我在这里同意（并赞成）@jterrace，但我添加了一个非常细微的变化，它返回装饰函数，并包含装饰器和装饰函数的参数。

class Decon:
    def __init__(self, a=None, b=None, c=True):
        self.a = a
        self.b = b
        self.c = c

    def __enter__(self):
        # only need to return self 
        # if you want access to it
        # inside the context
        return self 

    def __exit__(self, exit_type, exit_value, exit_traceback):
        # clean up anything you need to
        # otherwise, nothing much more here
        pass

    def __call__(self, func):
        def decorator(*args, **kwargs):
            with self:
                return func(*args, **kwargs)
        return decorator

【讨论】：

看起来不错！当世界开始正确更新到 3.6+ 时，我会坚持使用 docs.python.org/3/library/…，但虽然 2- 存在风险，但您的解决方案看起来非常诱人！

工作就像一个魅力。