在Python中同时作为装饰器和上下文管理器的函数？

Question

这可能有点超出范围，但主要是出于好奇……

有没有可能创建一个可以被调用的对象（函数或类），它同时充当上下文管理器和装饰器呢：

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')

Answer 1

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()

Answer 2

在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

Answer 3

从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__()这两个方法。