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

Solution 1

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

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

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

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

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

View more solutions

Comments

[Sven Marnach]

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.")

[jterrace]

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.

[Warz]

How would you pass arguments to the decorator if the context manager takes arguments ?

[guettli]

You can use contextlib2 if you must use Python2: contextlib2.readthedocs.org/en/latest

[VMAtm]

@Warz context manager already being created during Decontext(some_context_manager). __enter__ and __exit__ methods are pre-defined.

[Federico]

This might help:coderwall.com/p/0lk6jg/…

[Delgan]

Related question using Decontext with arguments: Passing arguments to decontext decorator

[Jacob Oscarson]

Looks nice! When the world starts to be properly updated to 3.6+ I'll stick to docs.python.org/3/library/… but while 2- is a risk your solution looks very tempting!

[Jacob Oscarson]

Worked lika a charm.

[Guimoute]

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?

[Sven Marnach]

@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.

[coler-j]

If example_manager yielded a result, how would we get access to that when used as a decorator?

[mudassirkhan19]

@coler-j did you get an answer for it?