djangoproject.com | python.org | linux.com
demongin.org - Python: Decorate a Function Without Losing <i>self</i>

Python: Decorate a Function Without Losing self

In which I lay to rest a long-standing head-scratcher involving decorated unittest classes in Selenium test scripts that lose their original class attributes once wrapped in a decorator.


Thursday, 2009-11-19 | Programming, Testing

"I invented 'It's a good thing' before you were even born. "
Martha Stewart

So, I solved an odd/obscure python problem recently, and I believe that posterity might actually find the story of how I did it beneficial.

First, the setup. I have a bunch of unittest scripts that execute Selenium routines. The unittest parts are fairly plain-Jane, insofar as they contain a unittest.TestCase class which contains functions named according to the "test_" convention. Here's an example of a very simple one:

class CompareInitialScreens(unittest.TestCase):
    """Use Selenium to get a screen capture and store it."""

    def test_compare_screens(self):   
        """MD5 sums must match for this test to pass."""

        self.sel = Browser()

        self.sel.setUp()
        self.sel.get_current_screen()
        self.sel.tearDown()
Super simple.

My problem was that I was running these things in batches (big batches) and I wanted them to spit out more than just the usual unittest output (which is fairly uninformative and non-specific, probably on purpose). Basically, what I had to do was have them say their class and function names on stderr in such a way as not to obstruct/obscure the natural output of unittest.

My original (very poor, brute-force, hack) solution was to do a bunch of fancy crap within each and every function. This was not only a huge violation of DRY, but it also really hurt the legibility of my scripts. Here's an example of what I was doing:
class CompareInitialScreens(unittest.TestCase):  
    """Use Selenium to get a screen capture and store it."""

    def test_compare_screens(self):  
        """MD5 sums must match for this test to pass."""
        test_class = self.__class__.__name__
        test_name = sys._getframe().f_code.co_name
        sys.stderr.write("\n" + repr(self.__doc__).strip("''"))
        sys.stderr.write("\nResults for '%s' (%s)\n" % (test_name, test_class))

        self.sel = Browser()

        self.sel.setUp()
        self.sel.get_current_screen()
        self.sel.tearDown()
See all of that junk that follows the doc string for the "test_compare_screens" function? I was using those same four lines over and over in a number of different scripts.

Now, I knew that what I was doing was wrong. I even knew that the solution to my problem was to use some kind of decorator function*, but I couldn't figure out how to make a decorator do what I wanted. The problem I kept having was that my functions would "forget" their original class once I tacked the decorator function onto them and start acting like the only class that they belonged to was the decorator class.

For instance, with code like this:
def __init__(self, f):
    print "DECORATOR INITIALIZED!"
    f() # Prove that function definition has completed

class TestRandomStuff(unittest.TestCase):
    seq = range(10)

    @Deco
    def testshuffle(self):
        # make sure the shuffled sequence does not lose any elements
        random.shuffle(self.seq)
I would repeatedly get traces like this:
toconnell@esme:~/sandbox-PYTHON$ ./deco_test.py 
DECORATOR INITIALIZED!
Traceback (most recent call last):
  File "./deco_test.py", line 15, in <module>
    class TestRandomStuff(unittest.TestCase):
  File "./deco_test.py", line 18, in TestRandomStuff
    @Deco
  File "./deco_test.py", line 9, in __init__
    f() # Prove that function definition has completed
TypeError: testshuffle() takes exactly 1 argument (0 given)</module>
The solution to my problem, I found out weeks later during a lull at work, was that I needed to use the wraps function from the functools module to explicitly preserve the original attributes of my wrapped functions.

What I ended up doing was writing a function which contains a slightly modified version of my four lines that I was putting in every function into a helper module that exists in a separate file from my test scripts:
def log_decorator(func):
    @wraps(func)
    def with_logging(*args, **kwargs):
        test_class = args[0].__class__.__name__
        test_name = func.__name__
        sys.stderr.write("\n" + repr(args[0].__doc__).strip("''"))
        sys.stderr.write("\nResults for '%s' (%s)\n" % (test_name, test_class))
        return func(*args, **kwargs)
    return with_logging
Then, in my individual test scripts, I would just import that little guy and wrap all of my test functions with him:
import helper
LogDecorator = helper.log_decorator

...

class CompareInitialScreens(unittest.TestCase): 
    """Use Selenium to get a screen capture and save it."""

    @LogDecorator   
    def test_compare_screens(self):   
        """MD5 sums must match for this test to pass."""

        self.sel = Browser()

        self.sel.setUp()
        self.sel.get_current_screen()
        self.sel.tearDown()
When I run the script with the decorator around my test functions, I get output like this:
toconnell@CRDL001:~/lemur$ ./tests/002a_compare_screens.py 

"Use Selenium to get a screen capture; compare MD5's with an existing file."
Results for 'test_compare_screens' (CompareInitialScreens)
.
OK

----------------------------------------------------------------------
Ran 1 test in 7.936s
Fuckin' sweet.



* This blog post is a great, practical crash course to decorators that is, IMO, a little bit better than Guido's official documentation: check it out.