Context manager as a decorator with access to the yielded object

Question

I have a context manager for an object that can be used similar to the open context manager, e.g.

with MyContextManager as cm:
    cm.do_something()

I know that a simple context manager can be made into a decorator if using contextlib.ContextDecorator to create it. Is it also possible to access the object yielded from the context manager if using a decorator? E.g. given the context manager above, something like:

@cmdecorator
def my_function(self, cm):
    cm.do_something

I couldn’t get that to work. Either I’m missing something trivial (hope so), or this is just not possible… It is only syntactic sugar in the end, but I’m interested if it is possible.

Asked By: MrBean Bremen

||

Source

Answer 1

No. This is explicitly mentioned in the documentation.

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.

Answered By: chepner

Answer 2

To answer my own question: while there is no automatic mechanism for this, it is easy enough to write a decorator that does exactly this (using a hardcoded keyword argument):

def patch_cm(f):
    @functools.wraps(f)
    def decorated(*args, **kwargs):
        with MyContextManager() as cm:
            kwds['cm'] = cm
            return f(*args, **kwargs)

    return decorated

which can be used (here in a unittest):

class MyContextManagerTest(TestCase):
    @patch_cm
    def test_context_decorator(self, cm):
        cm.do_something()
        self.assertTrue(cm.done())

which is the same as:

def test_context_decorator(self):
    with MyContextManager() as cm:
       cm.do_something()
       self.assertTrue(cm.done())

(what I actually use is a wrapper with arguments, but that is just one more wrapper around…)

This is also possible with a positional argument instead of a keyword argument (which I eventually ended up using):

def patch_cm(f):
    @functools.wraps(f)
    def decorated(*args, **kwargs):
        with MyContextManager() as cm:
            args = list(args)
            args.append(cm)
            return f(*args, **kwargs)

    return decorated

Answered By: MrBean Bremen

Context manager as a decorator with access to the yielded object

Question:

Answers: