Function acting as both decorator and context manager in Python?
PythonDecoratorContextmanagerPython Problem Overview
This might be pushing things a little too far, but mostly out of curiosity..
Would it be possible to have a callable object (function/class) that acts as both a Context Manager and a decorator at the same time:
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')
Python Solutions
Solution 1 - Python
Starting in Python 3.2, support for this is even included in the standard library. Deriving from the class contextlib.ContextDecorator
makes it easy to write classes that can be used as both, a decorator or a context manager. This functionality could be easily backported to Python 2.x -- here is a basic implementation:
class ContextDecorator(object):
def __call__(self, f):
@functools.wraps(f)
def decorated(*args, **kwds):
with self:
return f(*args, **kwds)
return decorated
Derive your context manager from this class and define the __enter__()
and __exit__()
methods as usual.
Solution 2 - Python
In Python 3.2+, you can define a context manager that is also a decorator using @contextlib.contextmanager
.
From the docs:
> contextmanager()
uses ContextDecorator
so the context managers it creates can be used as decorators as well as in with
statements
Example usage:
>>> 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
Solution 3 - Python
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
now you can do:
mydeco = Decontext(some_context_manager)
and that allows both
@mydeco
def foo(...):
do_bar()
foo(...)
and
with mydeco:
do_bar()
Solution 4 - Python
Here's an example:
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()
This prints:
init 1 2
enter
with
exit
init 3 4
call
about to call
sample
done calling
Solution 5 - Python
Although I agree (and upvoted) @jterrace here, I'm adding a very slight variation that returns the decorated function, and includes arguments for both the decorator and the decorated function.
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