Introduction to Python Decorators

I faced a situation recently while building a Web application where a Web service that we made requests from would occasionally hang when we made a request. Because we run our application on Linux, we decided to use signals to add a timeout to our request. Because we use Python, we implemented the timeout as a decorator, since we knew we would likely find other places where we needed to add timeouts. A simplified version of this code makes an excellent introduction to how decorators work.

Using signals

First, we should cover how to use signals. Listing 1 shows how to interrupt a runaway function using the alarm function from the signal module.

Listing one: using signals to interrupt runaway functions.

import logging, random, signal

logging.basicConfig(level=logging.DEBUG,
                    format="%(levelname)s %(asctime)s %(message)s")

class TimeoutException(Exception):
    "Indicates that the function has taken too long."

def handle_alarm(*args):
    logging.error("Timeout!")
    raise TimeoutException(*args)

def returns_sometimes():
    "50/50 chance -- return immediately or hang forever."

    if random.choice([False, True]):
        logging.debug("Hurray!")
        return 

    else:
        logging.debug("uh oh...")
        i = 1
        while "forever":
            i += 1
            if i == 1000000:
                logging.debug("another million iterations...")
                i = 0

def main():
    "Show how to cancel a timeout if the function does return in time."

    # Set the handler for the alarm signal.
    signal.signal(signal.SIGALRM, handle_alarm)

    # Schedule an alarm for two seconds in the future.
    signal.alarm(2)

    try:
        logging.debug("About to call the function that returns sometimes...")
        returns_sometimes()

        # Now cancel the impending alarm!
        signal.alarm(0)

        logging.debug("We got a response back!")

    except TimeoutException:
        logging.error("We didn't get a response back in time.")

if __name__ == "__main__": main()

Throughout this article, we’ll use a place-holder for the real function that connected to the Web service. The returns_sometimes function, defined on line 17, will either return immediately or hang indefinitely, depending on the outcome of the random choice between True and False.

The strategy we use to interrupt the runaway call is very simple. We set up a handler for the SIGALRM signal, so that when the signal arrives, we will execute the function handle_alarm. Then we schedule an alarm signal. Note that if our code didn’t handle the signal, the OS would kill our process.

The handle_alarm function logs the signal’s arrival and then raises a TimeoutException which is an exception that we define in this module. The handle_alarm function accepts* args as parameters and then adds those into the TimeoutException. In this code, we don’t need any of the information passed in, like a reference to the calling frame, but it’s available if a need for it turns up later.

We schedule the alarm after setting up the handler, rather than the other way around, because once the alarm is scheduled, the clock is ticking. You don’t want to waste milliseconds on extra function calls if they can be avoided by calling them before starting the alarm.

Next we enter a try / except block, do some logging, and then call our function returns_sometimes. At this point, our program blocks until that function returns. If it returns immediately, then the very first thing we do is call signal.alarm again, with a value of zero for the seconds. Any call to signal.alarm will delete the previous alarm scheduled. Calling signal.alarm with a zero will not schedule a new alarm.

After logging the successful return, the main function completes.

However, in the event that returns_sometimes begins the while” forever” loop on line 27, our program hangs indefinitely at that point, until finally the OS sends the SIGALRM signal to our process. Then we jump into the handle_alarm function.

That function then logs the runaway problem and then raises the TimeoutException.

The snippet below shows the results of an unlucky run of listing1.py, where the function hangs, followed by an instance where it returned immediately:

$ python listing1.py
DEBUG 2007-09-24 00:04:19,770 About to call the function that returns sometimes...
DEBUG 2007-09-24 00:04:19,770 uh oh...
DEBUG 2007-09-24 00:04:20,377 another million iterations...
DEBUG 2007-09-24 00:04:21,572 another million iterations...
ERROR 2007-09-24 00:04:21,784 Timeout!
ERROR 2007-09-24 00:04:21,784 We didn't get a response back in time.

$ python listing1.py
DEBUG 2007-09-24 00:04:34,116 About to call the function that returns sometimes...
DEBUG 2007-09-24 00:04:34,116 Hurray!
DEBUG 2007-09-24 00:04:34,116 We got a response back!

In the first call, the two-second interrupt is made clear by looking at the differences in the logging timestamps.

It is important to bear in mind that if we didn’t raise an exception from within our signal handler, when the handler finished, we would drop right back into the while” forever” loop.

See what happens after line 15 is commented out, so that we don’t raise the TimeoutException. The handle_alarm function executes, then returns, and we return immediately to the while loop.

$ python listing1.py
DEBUG 2007-09-23 23:59:12,461 About to call the function that returns sometimes...
DEBUG 2007-09-23 23:59:12,462 uh oh...
DEBUG 2007-09-23 23:59:14,159 another million iterations...
ERROR 2007-09-23 23:59:14,510 Timeout!
DEBUG 2007-09-23 23:59:14,756 another million iterations...
Traceback (most recent call last):
  File "a.py", line 55, in <module>
    if __name__ == "__main__": main()
  File "a.py", line 45, in main
    returns_sometimes()
  File "a.py", line 31, in returns_sometimes
    i = 0
KeyboardInterrupt

I killed the process by hitting Ctrl-C. Hitting Ctrl-C sends a different signal (SIGINT) to the process. Incidentally, you can handle a SIGINT signal in just the same way as how we handled SIGALRM.

A Few Simple Decorator Examples

The solution in Listing 1 is acceptable for the immediate problem at hand. However, the code in Listing 1 is not easy to reuse.

Fortunately, Python decorators allow the developer to wrap one function inside of another, so the results from the inner function can be” decorated” by the outer function. Decorators can be confusing at first, so we’ll start with some simple examples in Listing 2.

Listing two: simple demonstrations of decorators.

import logging

logging.basicConfig(level=logging.DEBUG,
                    format="%(levelname)s %(message)s")

def dec1(f):
    "Adds 1 to the returned value from the decorated function."

    def _f(*args, **kwargs):
        logging.debug("Inside decorated.")
        y = f(*args, **kwargs)
        return y + 1

    logging.debug("Finished building the decorated function.")

    return _f

@dec1
def square(x):
    return x * x

def dec2(N):

    """
    Return a decorator that adds N to the returned value, where N is a parameter
    passed into the function.
    """

    logging.debug("N is set to %d." % N)

    def decN(f):

        def _f(*args, **kwargs):
            logging.debug("Inside decorated.")
            y = f(*args, **kwargs)
            return y + N

        logging.debug("Finished building the decorated function.")

        return _f

    return decN

@dec2(7)
def cube(x):
    return x * x * x

# These examples show a decorator can be applied to a function after it is created.
def add_two_numbers(a, b): return a + b

add_six = dec2(6)

add_two_numbers_plus_six = add_six(6)

add_five_to_the_sum_of_a_and_b = dec2(5)(add_two_numbers)

def main():

    logging.info("results from square(3): %s." % square(3))
    logging.info("results from cube(2): %s." % cube(2))

    logging.info("results from add_five_to_the_sum_of_a_and_b(2, 3): %d." %
                 (add_five_to_the_sum_of_a_and_b(2, 3)))

if __name__ == "__main__": main()

Note that the square function we define starting on line 24 returns the square of its parameter, but when we call it, it returns something slightly different:

$ python
>>> import listing2
>>> listing2.square(3)
DEBUG Inside decorated.
10

This is because when we defined square we added the line @dec1 above. This means that what we think of as the function square is in fact the function square after it has been decorated by the dec1 decorator. The @dec1 syntax is syntactic sugar for this operation instead:

def square(x):
    return x * x

square = dec1(square)

In other words, we first define a function square, then re-bind the symbol square to the function returned by dec1 when we pass in square as a parameter. In many instances, I apply the decorator after the fact, because I am importing a function from a previously-written module.

Additionally, there are plenty of times when I don’t want to permanently attach my decorator. Defining a decorator is just like defining any other python function. They accept one parameter which is the function to decorate.

Then inside the decorator function, we build a new function that shows how to call the function we passed in. Inside dec1, for example, we define _f, and we give it the parameters *args and **kargs which allows the function to accept any, or no, parameters.

Inside _f, we call the original function f and store the result into the variable y. Then we return y+1, which is why when we called square(3), we got back to 10. It is important to keep straight that when we call dec1(square) or @dec1, the inner function _f does not execute. It is defined and returned only. All of the logging statements should make that fact clear:

>>> def f(a, b): return a + b
...
>>> f(1, 2)
3
>>> g = listing2.dec1(f)
DEBUG Finished building the decorated function.
>>> g(1, 2)
DEBUG Inside decorated.
4

Creating function g by decorating f didn’t execute the original function f. Instead we return a new function and store it in g. Only when we call g did the original f execute.

While dec1 was fairly straightforward, dec2 is slightly more sophisticated because it’s a decorator factory. In other words, we can pass parameters to dec2 to change how it decorates. In this case, dec2 accepts a parameter N and then returns a decorator that adds N to a function.

So, to emulate the dec1 decorator, we can do:

@dec2(1)

Writing a decorator factory means that instead of returning a function that is decorated, we return a function that can decorate another.

You may find that when writing decorator factories, it’s useful to start with by defining the end decorator, and then later define another decorator to build that, after replacing all the constants with parameters. For example, I knew I wanted to build dec2 that would accept a parameter and let me decorate any function adding that parameter to it. So, first I wrote dec1.

Then when I stared work on dec2, I copied the definition of dec1 inside of dec2 and indented it one level and renamed dec1 to decN. Then I changed line 42 so that instead of returning y+1, the function returns y+N.

Implementing Timeouts with Decorators

Listing 3 shows we implement the same timeout logic from Listing 1, but this time, as a decorator. The with_timeout decorator is similar to dec2 from Listing 3 in that it is really a decorator factory. We pass in an argument for how long to wait, and then another argument that is a function that we’ll execute if the alarm goes off.

Listing three: implement listing 1 as a decorator.

import logging, signal

from listing1 import TimeoutException, handle_alarm, returns_sometimes

logging.basicConfig(level=logging.DEBUG,
                    format="%(levelname)s %(message)s")

def with_timeout(seconds_to_wait, on_timeout):

    """
    After seconds_to_wait seconds, we interrupt the function and instead call
    on_timeout.
    """

    def dec(f):

        def _f(*args, **kwargs):

            # Set the handler for the alarm signal.
            signal.signal(signal.SIGALRM, handle_alarm)

            # Schedule an alarm for two seconds in the future.
            signal.alarm(seconds_to_wait)

            try:
                # Call the function.
                y = f(*args, **kwargs)

                # If the function returned, then cancel the alarm.
                signal.alarm(0)

                # Now return what we got.
                return y 

            except TimeoutException, ex: on_timeout(ex)

        return _f

    return dec

def out_of_time(ex):

    "Clean up and move on."

    logging.debug("We got unlucky.")
    logging.debug("ex.args is %s." % str(ex.args))

quit_after_2  = with_timeout(2, out_of_time)(returns_sometimes)

def main():

    logging.debug("Calling returns_sometimes after applying with_timeout decorator.")

    quit_after_2()

    logging.debug("We returned!")

if __name__ == "__main__": main()

We ask for a function on_timeout that will be called if the function times out rather than raising the exception directly so that the underlying function doesn’t have to specify how to handle timeout. Remember– we’re trying to keep the timeout logic separate from what the function is supposed to do.

Preserving the Function Signature

One negative consequence of applying decorators to functions is that the original signature is destroyed. The function returns_sometimes defined in Listing 1 takes no parameters. But the function quit_after_2, which is really just the exact same function inside a decorator, appears to accept any parameters.

Using the inspect module from the standard library shows the problem:

>>> inspect.getargspec(listing1.returns_sometimes)
([], None, None, None)
>>> inspect.getargspec(listing3.quit_after_2)
([], 'args', 'kwargs', None)

If you’re not familiar with inspect.getargspec, it returns a tuple that describes the arguments that the given function accepts. The tuple has four elements– the first element is a list of arguments; the second element is either None or the name of the* argument; the third element is the name of** argument; and finally, the last element is either None or is a list of any arguments with default values.

The reason why their signatures is different is that inside of with_timeout, we defined the function _f that accepts the parameters *args and **kwargs. This may seem like a minor point, but dynamic languages like Python that allow for runtime introspection can do some really amazing things.

Fortunately, Michele Simionato has solved this problem for us with his decorator module available here.

Listing four: use Michele Simionato’s decorator module to preserve the function signature.

import inspect, logging, signal

from decorator import decorator

import listing1, listing3

logging.basicConfig(level=logging.DEBUG,
                    format="%(levelname)s %(message)s")

def better_with_timeout(seconds_to_wait, on_timeout):

    """
    After seconds_to_wait seconds, we interrupt the function and instead call
    on_timeout.
    """

    def dec(f, *args, **kwargs):

        # Set the handler for the alarm signal.
        signal.signal(signal.SIGALRM, listing1.handle_alarm)

        # Schedule an alarm for two seconds in the future.
        signal.alarm(seconds_to_wait)

        try:
            # Call the function.
            y = f(*args, **kwargs)

            # If the function returned, then cancel the alarm.
            signal.alarm(0)

            # Now return what we got.
            return y 

        except listing1.TimeoutException, ex: on_timeout(ex)

    return dec

two_second_interrupt = better_with_timeout(3, listing3.out_of_time)

better_quit_after_2 = decorator(two_second_interrupt)(listing1.returns_sometimes)

@decorator(two_second_interrupt)
def z():
    listing1.returns_sometimes()

def main():

    logging.debug("inspect.getargspec(listing3.quit_after_2): %s"
                  % str(inspect.getargspec(listing3.quit_after_2)))

    logging.debug("inspect.getargspec(better_quit_after_2): %s"
                  % str(inspect.getargspec(better_quit_after_2)))

    better_quit_after_2()

    z()

if __name__ == "__main__": main()

Listing 4 shows a better implementation of the with_timeout decorator, unimaginatively named better_with_timeout that depends on the decorator module. Running the main method from Listing 4 shows how we now get accurate information when we inspect the functions decorated with the decorator module.

>>>; listing4.main()
DEBUG 2007-09-24 22:02:44,206 inspect.getargspec(listing3.quit_after_2): ([], 'args', 'kwargs', None)
DEBUG 2007-09-24 22:02:44,207 inspect.getargspec(better_quit_after_2): ([], None, None, None)

Conclusion

Now we have a decorator than can be applied to any arbitrary function to add completely new functionality, either when the function is defined, or after the fact. More importantly, I hope I illustrated some of the beauty and power that Python offers.