3

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?

clemens
  • 16,716
  • 11
  • 50
  • 65
ocwirk
  • 1,079
  • 1
  • 15
  • 35

1 Answers1

7

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()
Robᵩ
  • 163,533
  • 20
  • 239
  • 308