How to use a context manager inside a decorator and how to pass an object created in decorator to decorated function

Question

I have a test class which requires to do some cleanup at the end. To make sure that the user won't forget to do this, I want to add a context manager to the class. I also have a decorator, inside which I want to use this context manager to create an object of test class and pass it to the decorated function. Is it even possible?

This is what I am looking to do:

class test:
    def __init__(self, name):
        self._name = name
        print "my name is {0}".format(name)

    def exit():
        print "exiting"

    @contextmanager
    def testcm(self):
        print "inside cm"
        try:
            yield self
        finally:
            self.exit()

    def randomtest(self, str):
        print "Inside random test {0}".format(str)


def decorate(name):
    def wrapper(testf):
        def testf_wrapper(test):
            with test(name).testcm() as testobj:
                return testf(testobj)
            return testf_wrapper
        return wrapper
    return decorate

@decorate("whatever")
def testf(testobj):
    testobj.randomtest("randomness")

The function testf takes the test class object - testobj and do stuff with it. Afterwards, because of the context manager, testcm makes sure the cleanup function is called.

So there are two questions:

1. How do I use a context manager inside a decorator, from what I know a decorator must return a function, but if I return the function(as in code above), how would context manager going to call cleanup?

2. How do I pass an object created in decorator to the decorated function, if I pass it like in above code, how would I call the decorated function?

Answer 1

Your example program has several errors in it, and I got lost in all of the test/testf/testobj redundancy. Allow me to address your questions directly.

How do I use a context manager inside a decorator?

Exactly as you would use a context manager anywhere else. Consider this program, which uses a decorator to transparently convert str to file while invoking a function:

def opener(func):
    def wrapper(name):
        with open(name) as input_file:
            func(input_file)
    return wrapper

@opener
def first_line(fd):
    print fd.readline()

first_line('/etc/passwd')

As you can see, the decorator function uses a context manager around the invocation of the decorated function.

How do I pass an object created in decorator to the decorated function, if I pass it like in above code, how would I call the decorated function?

Exactly as you would pass an object to any function. See my example above. The decorator creates a file object and passes it to the decorated function.

---

For completeness, here is your sample program with the errors fixed:

from contextlib import contextmanager

class test:
    def __init__(self, name):
        self._name = name
        print "my name is {0}".format(name)

    def exit(self):
        print "exiting"

    @contextmanager
    def testcm(self):
        print "inside cm"
        try:
            yield self
        finally:
            self.exit()

    def randomtest(self, str):
        print "Inside random test {0}".format(str)


def decorate(name):
    def wrapper(testf):
        def testf_wrapper():
            with test(name).testcm() as testobj:
                return testf(testobj)
        return testf_wrapper
    return wrapper

@decorate("whatever")
def testf(testobj):
    testobj.randomtest("randomness")


testf()