Signature Preserving Function Decorators

(The code examples in this post are also available here.)

On the comments to the post on the greatest common divisor, Steve pointed out that the code to time function performance could be embedded in the code in a neater way using function decorators. At the time, I confess I did not have a clue of what a decorator is. But I’ve been doing my homework, and so here’s my take on the subject…

So what the @#$%& is a decorator?

Many moons ago, when I ditched my ZX Spectrum for a 80286 PC (with a i287 math coprocessor), and moved from BASIC to C, I eventually found out that there were such things as pointers to functions. What a sick thing to do to functions! What were this Kernighan and Ritchie!? Perverts!? So, still mourning for my lost GOTOs and GOSUBs, I dismissed the whole subject as pornographic, and decided not to learn it on moral grounds. Well, I’ve come a long way since then, am much more open-minded now, and have even developed a keen appreciation for errr… adult entertainment. So, bring in the syntactic porn sugar!!!

The basic idea is that when you write something like…

@my_decorator
def decorated_function(a, b, c) :
    <function definition goes here>

…it is evaluated the same as if you had written…

def decorated_function(a, b, c) :
    <function definition goes here>
decorated_function = my_decorator(decorated_function)

If this last assignment has sent you into stupor, that is simply the sweet way in which python’s indirection handles function pointers, so all there is to it is that my_decorator has to be a function (or a callable object, i.e. one with a __call__ method) that takes a function as only argument, and returns a function (or a callable object). This last return function has to be compatible with the arguments of decorated_function, since calls to the latter will actually call the former.

You can even raise the bar by passing arguments to the decorator, but the basic scheme is the same, if you write this…

@my_decorator(dec_a, dec_b, dec_c)
def decorated_function(a, b, c) :
    <function definition goes here>

…python does the same as if you had written this…

def decorated_function(a, b, c) :
    <function definition goes here>
decorated_function = my_decorator(dec_a, dec_b, dec_c)(decorated_function)

…where my_decorator(dec_a, dec_b, dec_c) should return a function, possibly different depending on the passed parameters, which takes a function as a parameter, which should return yet another function, that better be compatible with the arguments of decorated_function, for the same reasons as before.

Decorators can even be nested, but the functioning is quite predictable s well. You write this…

@one_decorator(one_args)
@two_decorators(two_args)
def decorated_function(a, b, c) :
    <function definition goes here>

…and python won’t know the difference from doing this…

def decorated_function(a, b, c) :
    <function definition goes here>
decorated_function = one_decorator(one_args)(two_decorator(two_args)(decorated_function))

…and by now you should be able to figure out what is being returned by what.

If you are still in distress, it may be a good idea to stop here and first read David‘s great article at Siafoo , or Bruce Eckel‘s series (1, 2, 3) at Artima. To see them suckers at work, check the Python Decorator Library. There’s also the option to drink from the pristine springs of PEP 318, not that I recommend it to the uninitiated, though.

A Decorator to Time Function Execution

So lets learn by doing, and build a decorator to time the performance of the decorated function, trying to mimic as much as possible the approach I have been following up till now: the last of the function’s argument is verbose, which defaults to False, and setting it to True activates printout of the time required by the function to run. This is possible, but messy. On the other hand, what can more easily be achieved is to control the timing with a keyword argument, as we will implement…

I have added more goodies… I learnt from the folks at Leohnard Euler’s Flying Circus that, while time.clock has superior performance than time.time under Windows, it is the opposite on Linux, so we will choose our timing function at run time, through os.name, which holds a string with the operating system your computer is running.

from __future__ import division
from os import name as OP_SYS
import time

def performance_timer_1(_func_) :
    """
    A decorator function to time execution of a function.

    Timing is activated by setting the verbose keyword to True.
    """

    stopwatch = time.time
    if OP_SYS in ('nt', 'dos') :
        stopwatch = time.clock

    def _wrap_(*args, **kwargs) :
        """
        This is a function wrapper.
        """
        
        verbose = kwargs.pop('verbose',False)
        t = stopwatch()
        ret = _func_(*args, **kwargs)
        t = stopwatch() - t
        if verbose :
            print "Call to",_func_.__name__,"took",t,"sec."
        return ret

    return _wrap

Run this file on IDLE, or import it into an interactive python session, and you can do things like these:

>>> @performance_timer_1
def decorated_sum(a,b) :
"""
This is a function.
"""
return a+b

>>> decorated_sum(3,5)
8

>>> decorated_sum(3,5,verbose=True)
Running function decorated_sum took 2.72380986743e-06 sec.
8

If you are wondering about the silly docstrings in the function and the wrapper, try the following:

>>> help(decorated_sum)
Help on function wrapper:

wrapper(*args, **kwargs)
This is a function wrapper.

Now that’s odd, isn’t it? Well, not really, as it is a natural byproduct of the devilish practice of assigning functions to functions: we are getting help for the function with which the decorator replaced the original function during decoration. Part of it is easy to fix, though, as the following code demonstrates…

import inspect

def performance_timer_2(_func_) :
    """
    A decorator function to time execution of a function.

    Timing is activated by setting the verbose keyword to True.
    """

    stopwatch = time.time
    if OP_SYS in ('nt', 'dos') :
        stopwatch = time.clock

    def _wrap_(*args, **kwargs) :
        """
        This is a function wrapper.
        """
        
        verbose = kwargs.pop('verbose',False)
        t = stopwatch()
        ret = _func_(*args, **kwargs)
        t = stopwatch() - t
        if verbose :
            print "Call to",_func_.__name__,"took",t,"sec."
        return ret
    _wrap_.__name__ = _func_.__name__
    _wrap_.__doc__ = _func_.__doc__
    return _wrap_

If you try the same thing again, this time you will get…

@performance_timer_2
def decorated_sum(a,b) :
"""
This is a function.
"""
return a+b

>>> help(decorated_sum)
Help on function decorated_sum:

decorated_sum(*args, **kwargs)
This is a function.

We’re almost there!!! There is just one disturbing thing remaining: the function signature. See, if we had invoked help on the undecorated decorated_sum, we would have got decorated_sum(a, b), not the generic arguments and keyword arguments thing.

Preserving the Function Signature

There are ways to get around it, but things are going to get messy. The primal source is Michele Simionato‘s decorator module. Most of what I will present in the rest of this post is shamelessly ripped off from inspired on his code. You can check the documentation, or the last section of David’s Siafoo article, if you want to use that module.

There are two problems though… The first is of course having to rely on a third party module which isn’t part of the standard library: the rules forbid it!

But there is also a more fundamental issue: we would like to preserve the original function signature, but at the same time we are changing it!!! Going back to the prior example, if we preserved the signature unaltered as decorated_sum(a, b), we would get a TypeError: decorated_sum() got an unexpected keyword argument 'verbose' every time we tried to activate the timer. Trust me, I found the hard way…

What we would like to get is something like decorated_sum(a, b, **kwargs). To achieve something like this we are going to have to dynamically generate a function at run-time…

def performance_timer_3(_func_) :
    """
    A decorator function to time execution of a function.

    Timing is activated by setting the verbose keyword to True.
    """
    
    stopwatch = time.time
    if OP_SYS in ('nt', 'dos') :
        stopwatch = time.clock

    def _wrap_(*args, **kwargs) :
        """
        This is a function wrapper.
        """
        
        verbose = kwargs.pop('verbose',False)
        t = stopwatch()
        ret = _func_(*args, **kwargs)
        t = stopwatch() - t
        if verbose :
            print "Call to",_func_.__name__,"took",t,"sec."
        return ret
        
    sig = list(inspect.getargspec(_func_))
    wrap_sig = list(inspect.getargspec(_wrap_))
    if not sig[2] :
        sig[2] = wrap_sig[2]
    src =  'def %s%s :\n' %(_func_.__name__, inspect.formatargspec(*sig))
    src += '    return _wrap_%s\n' % (inspect.formatargspec(*sig))
    evaldict = {'_wrap_' : _wrap_}
    exec src in evaldict
    ret = evaldict[_func_.__name__]
    ret.__doc__ = _func_.__doc__
    return ret

If you are not familiar with them, you may want to check the documentation for the inspect module, as well as for the exec statement, to get a hang of what’s going on. It is important to note how the _wrap_ function is not simply named in the dynamically generated source code, but included in evaldict. If we didn’t do it this way, it would be out of scope once we left the decorator, and an error will occur. Again, you can trust me on this one…

But the thing is we have finally arrived where we wanted to…

>>> @performance_timer_3
def decorated_sum(a,b) :
"""
This is a function.
"""
return a+b

>>> help(decorated_sum)
Help on function decorated_sum:

decorated_sum(a, b, **kwargs)
This is a function.

So every time you want to do a function decorator with a wrapper activated by a keyword, all you need to do is take this last example as a template, and with some cutting and pasting you’ll be done in no time.

Now wait a minute!!! Wasn’t avoiding this kind of thing what decorators where supposed to be all about? Didn’t we get into this mess to avoid cutting and pasting a couple of lines of code to the beginning and end of every function? What if we built a decorator that automated the decorating process?

def signature_decorator(_wrap_, has_kwargs = False) :
    """
    A decorator to create signature preserving wrappers.

    _wrap_     : the wrapper function, which must take a function as first
                 argument, can then take several optional arguments, which
                 must have defaults, plus the *args and **kwargs to pass to
                 the wrapped function.
    has_kwargs : should be set to True if the wrapper also accepts keyword
                 arguments to alter operation.
    """
    def decorator(func):
        """
        An intermediate decorator function
        """
        sig = list(inspect.getargspec(func))
        wrap_sig = list(inspect.getargspec(_wrap_))
        sig[0], wrap_sig[0] = sig[0] + wrap_sig[0][1:], wrap_sig[0] + sig[0]
        wrap_sig[0][0] = '_func_'
        sig[3] = list(sig[3]) if sig[3] is not None else []
        sig[3] += list(wrap_sig[3]) if wrap_sig[3] is not None else []
        sig[3] = tuple(sig[3]) if sig[3] else None
        wrap_sig[3] = None
        if sig[2] is None and has_kwargs:
            sig[2] = wrap_sig[2]
        wrap_sig[1:3] = sig[1:3]
        src = 'def %s%s :\n' % (func.__name__, inspect.formatargspec(*sig))
        src += '    return _wrap_%s\n' % (inspect.formatargspec(*wrap_sig))
        evaldict = {'_func_':func, '_wrap_': _wrap_}
        exec src in evaldict
        ret = evaldict[func.__name__]
        ret.__doc__ = func.__doc__
        return ret
    return decorator
    
def timer_wrapper_1(_func_, verbose = False, *args, **kwargs) :
    """
    A wrapper to time functions, activated by the last argument.

    For use with signature_decorator
    """
    stopwatch = time.time
    if OP_SYS in ('nt','dos') :
        stopwatch = time.clock
    t = stopwatch()
    ret = _func_(*args,**kwargs)
    t = stopwatch() - t
    if verbose :
        print "Call to",_func_.__name__,"took",t,"sec."
    return ret

def timer_wrapper_2(func, *args, **kwargs) :
    """
    A wrapper to time functions, activated by a keyword argument..

    For use with signature_decorator
    """
    stopwatch = time.time
    if OP_SYS in ('nt','dos') :
        stopwatch = time.clock
    verbose = kwargs.pop('verbose',False)
    t = stopwatch()
    ret = func(*args,**kwargs)
    t = stopwatch() - t
    if verbose :
        print "Call to",_func_.__name__,"took",t,"sec."
    return ret                                     

Do note that the wrappers are now defined outside the decorator, and that they must take the wrapped function as first argument. Do note as well the decorators two parameters, the first being the wrapper, the second whether the wrapper takes keyword arguments or not. To get a hang of how this works try the following…

>>> @signature_decorator(timer_wrapper_1)
def decorated_sum(a,b=2) :
"""
This is a function.
"""
return a+b

>>> help(decorated_sum)
Help on function decorated_sum:

decorated_sum(a, b=2, verbose=False)
This is a function.

>>> @signature_decorator(timer_wrapper_2,True)
def decorated_sum(a,b=2) :
"""
This is a function.
"""
return a+b

>>> help(decorated_sum)
Help on function decorated_sum:

decorated_sum(a, b=2, **kwargs)
This is a function.

There are a million things that could go wrong with this last decorator, specially with conflicting variable or keyword names in the function and the wrapper. If you check the decorator module source code, you can find ways to build robustness to such issues. As for myself, a week’s worth of sleep deprivation is enough: I already know more about decorators than ever thought possible or wished. If you are curious, you can check the final, a little more elaborate, version of the performance_timer decorator I’m sticking to, which does not use this last refinement of the decorator_signature, checking nrp_base.py.