Function acting as both decorator and context manager in Python?

Question

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

can you give an example of what it would do? are the two functionalities different? — jterrace, Feb 09 '12 at 15:24
It would basically being two alternate ways of setting up fixtures in test suites. — Jacob Oscarson, Feb 09 '12 at 15:31

Sven Marnach · Accepted Answer · 2012-02-09T15:39:36.403

74

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.

edited Feb 09 '12 at 15:39

answered Feb 09 '12 at 15:26

Sven Marnach

574,206
118
941
841

8

You can use contextlib2 if you must use Python2: http://contextlib2.readthedocs.org/en/latest/ – guettli Sep 02 '14 at 06:50
2

This might help:https://coderwall.com/p/0lk6jg/python-decorators-vs-context-managers-have-your-cake-and-eat-it – Federico Feb 24 '16 at 10:35
Using python 3.8.5, adding or removing the line `@functools.wraps(f)` does not change anything and both context manager and argument-less decorator work. Why is that? – Guimoute Jan 05 '21 at 15:07
1

@Guimoute Please read the documentation of `functools.wraps()` to understand what it does. It doesn't change the way the wrapper function _works_, but it changes the way it _looks_ when inspecting it. For example the wrapper inherits the docstring and the name of the wrapped function. – Sven Marnach Jan 05 '21 at 15:11

score 56 · Answer 2 · answered Aug 08 '16 at 17:55

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

If `example_manager` yielded a result, how would we get access to that when used as a decorator? — coler-j, Nov 17 '21 at 19:59
@coler-j You can't from https://docs.python.org/3/library/contextlib.html >Note that there is one additional limitation when using context managers as function decorators: there’s no way to access the return value of __enter__(). If that value is needed, then it is still necessary to use an explicit with statement. — Hacker, Sep 10 '22 at 09:29

score 17 · Answer 3 · edited Feb 09 '12 at 16:07

17

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

edited Feb 09 '12 at 16:07

Sven Marnach

574,206
118
941
841

answered Feb 09 '12 at 15:38

nosklo

217,122
57
293
297

3

How would you pass arguments to the decorator if the context manager takes arguments ? – Warz Aug 07 '14 at 17:54
@Warz context manager already being created during `Decontext(some_context_manager)`. `__enter__` and `__exit__` methods are pre-defined. – VMAtm Oct 24 '15 at 11:56
Related question using `Decontext` with arguments: [Passing arguments to decontext decorator](https://stackoverflow.com/questions/25190979/passing-arguments-to-decontext-decorator) – Delgan Nov 04 '17 at 18:18

score 9 · Answer 4 · edited Aug 27 '19 at 08:47

9

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

edited Aug 27 '19 at 08:47

Jacob Oscarson

6,363
1
36
46

answered Feb 09 '12 at 15:35

jterrace

64,866
22
157
202

4

When used as a decorator, this won't act in the same way as a context manager, which seems to be the OP's intention. (See the OP's comment -- "It would basically being two alternate ways of setting up fixtures in test suites.") – Sven Marnach Feb 09 '12 at 15:43
1

Sure, the point was to show how to have a single class do both. I would leave it up to the OP to customize it for intended functionality. – jterrace Feb 09 '12 at 16:04

score 6 · Answer 5 · answered Aug 27 '19 at 01:41

6

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

answered Aug 27 '19 at 01:41

openwonk

14,023
7
43
39

Looks nice! When the world starts to be properly updated to 3.6+ I'll stick to https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator but while 2- is a risk your solution looks very tempting! – Jacob Oscarson Aug 27 '19 at 08:44
Worked lika a charm. – Jacob Oscarson Aug 27 '19 at 08:58
This solution allows one to use the arguments of the wrapped function or method, which none of the other solutions allows. – auxsvr Mar 22 '23 at 20:11

Function acting as both decorator and context manager in Python?

5 Answers5

Linked