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:
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?
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?