Python decorators and examples

Decorators without arguments Link to heading

Decorators modify functions. They can also modify classes, but they’re mainly used with functions.

Decorators are just a way to simplify syntax. Here’s an example of a decorator at work.

@some_decorator
def add(x,y): return x + y

This is the same thing as

def add(x,y): return x + y
add = some_decorator(add)

Without the decorator, we’d call some_decorator on add and assign the answer back to add. Decorators let you do this with less boilerplate.

Decorators with arguments behave quite differently to decorators without arguments. Let’s look at the simpler case first: decorators without arguments.

Here is the basic structure of a decorator:

def some_decorator(f): 
    def inner(*args, **kwargs): 
        return f(*args, **kwargs)
    return inner

where

  • some_decorator is the decorator
  • f is the function that is decorated (e.g. add from above)
  • *args and **kwargs are arguments to f
  • inner is the return value from some_decorator(add) that gets assigned back to add.

The outside layer under some_decorator executes once when you decorate the function, and the inner layer executes every time you call the function. So if you decorate add with some_decorator, when you define the function the outer loop some_decorator will run, and every time you call add the inner function will run.

Here’s a nice diagram (credit)

Let’s look at some examples.

Set documentation of a function to a specific string Link to heading

Below we have a function changedoc, a function that changes documentation string of another function. Let’s see how to use it both as a decorator and by itself.

def changedoc(f1): 
    f1.__doc__ = "a new doc"
    return f1

# Here is changedoc used not as a decorator
def somefun(x): return x+1
somefun = changedoc(somefun)
somefun.__doc__  # returns "a new doc"

# Here is changedoc as a decorator - a bit neater
@changedoc
def anotherfun(x): return x+2
anotherfun.__doc__  # returns "a new doc"

Add a specific string to the end of a function Link to heading

Here is another example. The decorator add_hello adds the word “hello” to the end of a function. I’ll show examples of how to use this both as a decorator, and not as a decorator.

We use *args and **kwargs as parameters in the inner function. This is useful because it holds arguments and keyword arguments you pass to f. This lets you run the function f through `f(*args, **kwargs).

A quick note. I find printing locals() in a function call is useful for trying to debug decorators because you can see exactly what variables are present. I’ve included some of this output from locals() to help understand what’s going on.

def add_hello(f): 
    def inner(*args, **kwargs): 
        print("Variables present: " , locals())
        return f(*args, **kwargs) + "_hello" 
    return inner

# Without decorator 
def print_alphabet(): return "abcdefghijklmnopqrstuvwxyz"
f_temp = add_hello(print_alphabet)
print(f_temp())

# With decorator
@add_hello
def print_alphabet(): return "abcdefghijklmnopqrstuvwxyz"
print(print_alphabet())

# Without decorator, and using args, kwargs
def repeat_x(x, n=4):    return x*n
f_temp = add_hello(repeat_x)
print(f_temp("deception ", n=6))

# With decorator, and using args, kwargs 
@add_hello
def repeat_x(x, n=4):    return x*n
print(repeat_x("deception ", n=6))

>>>
Variables present:  {'args': (), 'kwargs': {}, 'f': <function print_alphabet at 0x110533400>}
abcdefghijklmnopqrstuvwxyz_hello
Variables present:  {'args': (), 'kwargs': {}, 'f': <function print_alphabet at 0x1105338c8>}
abcdefghijklmnopqrstuvwxyz_hello
Variables present:  {'args': ('deception ',), 'kwargs': {'n': 6}, 'f': <function repeat_x at 0x110533d08>}
deception deception deception deception deception deception _hello
Variables present:  {'args': ('deception ',), 'kwargs': {'n': 6}, 'f': <function repeat_x at 0x1105420d0>}
deception deception deception deception deception deception _hello
>>>

Time how long a function runs for Link to heading

The decorator time_it below is a decorator used for timing functions.

To recap what’s going on:

  • The decorated function add_stuff goes into the f argument for time_it.
  • The arguments x and y in add_stuff are caught by *args in wrapper and are accessed in a tuple in the args variable. If we specified x or y by keyword arguments, these would be caught by **kwargs instead of *args.

Here is time_it:

import time 
def time_it(f): 
    def inner(*args, **kwargs): 
        t = time.time()
        result = f(*args, **kwargs)
        print (f.__name__ + " takes " +  str(time.time() - t) + " seconds.")
        return result 
    return inner

@time_it
def add_stuff(x,y): return x+y
add_stuff(1,2)

>>>
add_stuff takes 2.86102294921875e-06 seconds.
3
>>>

Log the internal state of a function Link to heading

Here is a decorator used for logging. Pretend the print statement actually writes to a file, and you get the idea.

def add_logs(f): 
    def inner(*args, **kwargs): 
        # insert real logging code here 
        print('write', *args, "to file")
        return f(*args, **kwargs)
    return inner    

@add_logs
def add_things(x,y,z): return x+y+z
add_things(1,4,3)

>>>
write 1 4 3 to file
8
>>>

Set a time limit on how often a function can be called Link to heading

Say we wanted a function to be called no more than once every minute. We could get this result by using a decorator, called once_per_min below.

Inside the inner function, we have to use the python keyword nonlocal to access the variable calltime defined out of its scope. We can’t use global because calltime isn’t a global variable.

def once_per_min(f): 
    calltime = 0 
    def inner(*args, **kwargs): 
        nonlocal calltime
        gap = time.time() - calltime
        if gap < 60: 
            msg = "You're calling this function too often. Try again in " + \
                          str(60 - gap) + " seconds."
            raise Exception(msg)
        calltime = time.time()
        return f(*args, **kwargs)
    return inner

@once_per_min
def add_stuff(x,y,z): return x+y+z
add_stuff(1,2,3)

# if we try again too quickly, it gives an error
add_stuff(1,2,3)

>>>
---------------------------------------------------------------------------

Exception                                 Traceback (most recent call last)

<ipython-input-10-5f783a66ca97> in <module>
      1 # if we try again too quickly, it gives an error
----> 2 add_stuff(1,2,3)


<ipython-input-7-4d8a29e52b68> in inner(*args, **kwargs)
      7             msg = "You're calling this function too often. Try again in " + \
      8                           str(60 - gap) + " seconds."
----> 9             raise Exception(msg)
     10         calltime = time.time()
     11         return f(*args, **kwargs)


Exception: You're calling this function too often. Try again in 55.97848320007324 seconds.
>>>

Make a function print out useful debugging information Link to heading

The decorator debug_this provides some useful information about a function it decorates.

import inspect, time
def debug_this(f): 
    def inner(*args, **kwargs): 
        print("[trace] Arguments of f:", args)
        print("[trace] Keyword arguments of f:", kwargs)
        print("[trace] f.__dict__:", f.__dict__)
        print("[trace] f.__name__:", f.__name__)
        # Uncomment if you want source code of function 
        # print("[trace] Source of f: \n####\n", inspect.getsource(f), "####")
        print("[trace] Starting execution of f")
        t1 = time.time()
        result = f(*args, **kwargs)
        t2 = time.time()
        print("[trace] Finished execution of f")
        print("[trace] f took", t2-t1,"seconds to run.")
        print("[trace] f returned: ", result)
        return result
    return inner

@debug_this
def somefun(a,b,c): 
    """This function is a bit complex but doesn't do anything interesting"""
    d = a + b 
    e = b + c 
    f = 10
    for x in (a,b,c,d,e): 
        f += x 
    return f 

somefun(1,2,4)

>>>
[trace] Arguments of f: (1, 2, 4)
[trace] Keyword arguments of f: {}
[trace] f.__dict__: {}
[trace] f.__name__: somefun
[trace] Starting execution of f
[trace] Finished execution of f
[trace] f took 3.814697265625e-06 seconds to run.
[trace] f returned:  26

26
>>>

Cache the results of a function Link to heading

This can be useful if a function takes a while to run. The fancy name for this is memoization.

We’ll use a dict as the cache, but we have to be careful how we implement the decorator. The obvious approach of using *args as a key to the cache dict won’t work, since it doesn’t catch keyword arguments, and args might contain unhashable types like lists or sets. To get around this, we pickle the arguments and keyword arguments to create a bytestring, and then use this bytestring as the key to the cache.

import pickle
def cache_this(f):
    cache = dict()
    def inner(*args, **kwargs):
        bs = (pickle.dumps(args), pickle.dumps(kwargs))
        if bs not in cache:
            print("caching result")
            cache[bs] = f(*args, **kwargs)
        else: 
            print("using cached result")
        return cache[bs]
    return inner

import time 
import numpy as np
@cache_this
def add_iterable(x): 
    """some long-running function with non-hashable arguments"""
    time.sleep(2)
    return np.sum(list(x))
add_iterable([1,2,5])    

>>>
caching result
8
>>>

add_iterable([1,2,5])    

>>>
using cached result 
8
>>>

Decorators with arguments Link to heading

All the above decorators had no arguments. Decorators can be used with arguments, and it is sometimes very useful to do so.

But the structure of these decorators is different. They need three nested layers, not two. In other words, we need to include a ‘middle` function.

def some_decorator(n): 
    def middle(f): 
        def inner(*args, **kwargs): 
            return f(*args, **kwargs)
        return inner  
    return middle

What’s going on?

  • The outer layer some_decorator is the name of the decorator, here once_per_n. The parameter to the decorator is also the parameter of the outer layer, which here I’ve called n. The outer layer executes once, when you provide an argument to the decorator.
  • The middle layer middle has as parameter f, the decorated function. This middle layer executes once, when we decorate the function.
  • The inner layer inner has as arguments *args and **kwargs, the arguments to the decorated function f. This inner layer is executed any time the decorated function f is called.

Another nice diagram (credit)

Let’s see some applications.

Stop a function running more than every n seconds Link to heading

If we modified the @once_per_min decorator to take an argument n, we could specify the length of time. Instead of an upper limit of one function call every minute, we could have one function call every n seconds.

Say we created a once_per_n decorator. Here’s how it would act on functions. It holds that

@once_per_n(5)
def add(x,y): return x+y

is the same as

add = once_per_n(5)(add)

The once_per_n(5)function also returns a function, that is then called on add.

Here is code for the once_per_n decorator. It’s similar to before, but we have a middle layer as well.

def once_per_n(n): 
    def middle(f):
        calltime = 0 
        def inner(*args, **kwargs): 
            nonlocal calltime; 
            gap = time.time() - calltime 
            if gap < n: 
                msg = "You're calling this function too often. Try again in " + \
                          str(n - gap) + " seconds."
                raise Exception(msg)
            calltime = time.time()
            return f(*args, **kwargs)
        return inner  
    return middle 

@once_per_n(5)
def add(x,y): return x+y
add(5,1)
# don't call this too quickly!
add(65,15)

>>>
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-24-aa3cda8aca6c> in <module>
    1 # don't call this too quickly!
----> 2 add(65,15)

<ipython-input-22-03225049070e> in inner(*args, **kwargs)
    8                 msg = "You're calling this function too often. Try again in " + \
    9                           str(n - gap) + " seconds."
---> 10                 raise Exception(msg)
    11             calltime = time.time()
    12             return f(*args, **kwargs)

Exception: You're calling this function too often. Try again in 4.640860080718994 seconds.
>>>

Add any string to the output of a function Link to heading

Before we had a @add_hello decorator to add the word hello to the end of any function. But if we made this decorator take an argument, we can add any string we like. Call this new decorator add_word.

# A decorator to add a word to the output of a function 
def add_word(word): 
    def middle(f): 
        def inner(*args, **kwargs):
            return f(*args, **kwargs) + word
        return inner
    return middle 

@add_word("_hello")
def alphabet(): return 'abcdefghijklmnopqrstuvwxyz'
print(alphabet())

@add_word("oh boy!!!")
def repeat_x(x, n=2): return x*n 
print(repeat_x("deception ", n=4))

>>>
abcdefghijklmnopqrstuvwxyz_hello
deception deception deception deception oh boy!!!
>>>

Make a numerical function return modulo n Link to heading

In this example the decorator mod_n modifies a function returning a number to return that number modulo n.

def mod_n(n): 
    def middle(f): 
        def inner(*args, **kwargs):
            return f(*args, **kwargs) % n
        return inner 
    return middle 


@mod_n(7) 
def add(x,y): return x + y
print(add(5,1), add(5,2), add(5,3))

>>>
6 0 1
>>>

Run tests when a function is defined Link to heading

Decorators can be useful for unit testing. Here run_tests runs tests against a function the first time it is defined.

This is useful to test for mistakes in the function definition, right as you write it. Also if you change something later in definition of the function, the tests will run automatically, and you’ll see if the function output has changed or not.

def run_tests(tests): 
    def middle(f): 
        for params, result in tests: 
            if f(*params) == result: print("Test", *params, "passed.")
            else:                    print("Test", *params, "failed.")
        def inner(*args, **kwargs):
            return f(*args, **kwargs)
        return inner 
    return middle 

# test cases for adds_to_ten below 
tests_eq_10 = [
    [(1,2,4), False],
    [(1,2,7), True],
    [(10,0,0),True], 
    [(-10,10,10),True],
    [(4,0,7), False]
]
@run_tests(tests_eq_10)
def adds_to_ten(x,y,z): return True if x+y+z==10 else False 

# test cases for adds_to_odd below 
tests_odd = [
    [(1,3,4), False], 
    [(1,3,5), True], 
    [(0,0,0), False], 
    [(0,0,1), True]
]
@run_tests(tests_odd)
def adds_to_odd(x,y,z): return (x+y+z)%2==1

>>>
Test 1 2 4 passed.
Test 1 2 7 passed.
Test 10 0 0 passed.
Test -10 10 10 passed.
Test 4 0 7 passed.
Test 1 3 4 passed.
Test 1 3 5 passed.
Test 0 0 0 passed.
Test 0 0 1 passed.
>>>

Inherit documentation from another function Link to heading

Making a function inherit documentation from another function is a useful job for decorators. In this example our decorator is copy_docs, the function with documentation we want to copy f_with_docs, and the decorated function f.

Note that there is no inner function below, because we don’t need to ever execute f. We are just copying across documentation and leaving it otherwise unchanged.

def copy_docs(f_with_docs):
    def middle(f):
        f.__doc__  = f_with_docs.__doc__
        f.__name__ = f_with_docs.__name__
        return f
    return middle

import numpy as np
@copy_docs(np.add)
def add(x,y): return x+y
print(add.__doc__[0:300])  # same documentation as numpy add function 

>>>
add(x1, x2, /, out=None, *, where=True, casting='same_kind', order='K', dtype=None, subok=True[, signature, extobj])

Add arguments element-wise.

Parameters
----------
x1, x2 : array_like
    The arrays to be added.  If ``x1.shape != x2.shape``, they must be
    broadcastable to a common shape (whi
>>>

Decorators tend to lose the documentation and attributes of the original function, and it can be a problem If you apply a decorator to a function, the decorated function won’t keep its __name__ and __doc__ attributes: they’ll be replaced with the name and docstring of something like middle or inner above. Also lost is information about the arguments of the function, so if you look at the help docs the arguments will be something like (*args, **kwargs), instead of the original arguments. There are some other things that are also lost.

A solution to this common problem is functools.wraps, a decorator from functools that is used to maintain the name, docstring, and other things of the original function.

Use functools.wraps in the decorator definitions. Let’s demonstrate it with the modulo example again.

from functools import wraps 
def modulo_n(n): 
    def middle(f): 
        @wraps(f)# adding @wraps here keeps info about f
        def inner(*args, **kwargs): 
            return f(*args, **kwargs) % n 
        return inner 
    return middle 


# some function with documentation 
@modulo_n(5)
def add(x,y): 
    """Adds two numbers, x and y, to make a third number, and then return it. """
    return x+y 

add(5,54)
add.__doc__  # docstring retained 

>>>
'Adds two numbers, x and y, to make a third number, and then return it. '
>>>

You can also use @wraps to directly transfer documentation from one function to another:

@wraps(np.add)
def add(x,y): return x+y
print(add.__doc__[0:300])  # same documentation as numpy add function 

>>>
add(x1, x2, /, out=None, *, where=True, casting='same_kind', order='K', dtype=None, subok=True[, signature, extobj])

Add arguments element-wise.

Parameters
----------
x1, x2 : array_like
    The arrays to be added.  If ``x1.shape != x2.shape``, they must be
    broadcastable to a common shape (whi
>>>

Decorators as classes Link to heading

You don’t have to write a decorator as a function. You can also write it as a class: it just has to implement the __call__ method, meaning that the class is callable.

Here is the basic structure. Say we were decorating a function called add.

  • The decorated function gets passed in as the parameter f to __init__, which assigns it to self.f so that other methods can also access it.
  • f is called in __call__ and __call__ also gets passed the *args and **kwargs of the f too.
  • __init__ gets called when the decorated function add is defined, but not when it is run.
  • __call__ gets called when the decorated function add is run, but not when it was defined.

This below example works but doesn’t do anything. Look at the printed output to see the control flow of the decorator.

class some_decorator: 
    def __init__(self, f): 
        print('in init statement')
        self.f = f 
    
    def __call__(self, *args, **kwargs):
        print('in call statement')
        return self.f(*args, **kwargs)
    
@some_decorator
def add(x,y):  return x+y    

>>>
in init statement
>>>

add(1,5)

>>>
in call statement
6
>>>

Here is the same decorator in function form. Note that the outer layer is basically the same as __init__ and the inner layer same as __call__.

def some_decorator2(f): 
    print('equivalent of __init__ statement')
    def inner(*args, **kwargs): 
        print('equivalent of __call__ statement')
        return f(*args, **kwargs)
    return inner

@some_decorator2
def add(x,y):  return x+y    

>>>
equivalent of __init__ statement
>>>

add(1,4)

>>>
equivalent of __call__ statement
5
>>>

What about decorators with arguments? Let’s implement class-style @once_per_n, the “wait n seconds between function calls” decorator from above.

A few differences.

  • The __init__ function now takes parameters from the decorator. __init__ is called when the decorated function is defined.
  • The _call__ function is now nested, and f is now an argument to __call__. This __call__ function now gets called when the decorated function is defined , not when it is run.
  • The function that runs every time the decorated function is called is now in inner. Whenever the decorated function is run, inner is executed.

I’ve kept in a few logging messages to make it a bit easier to see what happens.

import time 
class once_per_n: 
    def __init__(self, n): 
        self.n = n
        print("inside init")
        
    def __call__(self, f):
        self.calltime = 0 
        print("inside call ")
        def inner(*args, **kwargs):
            print("inside inner")
            gap = time.time() - self.calltime 
            if gap < self.n:
                msg = "You're calling this function too often. Try again in " + \
                              str(self.n - gap) + " seconds."
                raise Exception(msg)
            self.calltime = time.time()       
            return f(*args, **kwargs)
        return inner 

@once_per_n(6)
def add(x,y): return x+y

>>>
inside init
inside call 
>>>

add(5,5)
add(5,5) # will be run too frequently 

>>>
inside inner
inside inner

---------------------------------------------------------------------------

Exception                                 Traceback (most recent call last)

<ipython-input-12-e20425f72f86> in <module>
      1 add(5,5)
----> 2 add(5,5) # will be run too frequently


<ipython-input-11-a151d7e46213> in inner(*args, **kwargs)
     14                 msg = "You're calling this function too often. Try again in " + \
     15                               str(self.n - gap) + " seconds."
---> 16                 raise Exception(msg)
     17             self.calltime = time.time()
     18             return f(*args, **kwargs)


Exception: You're calling this function too often. Try again in 5.999882936477661 seconds.
>>>

Let’s compare the class-style decorator above to the function-style decorator below. Some points:

  • The outer layer once_per_n corresponds with __init__ above. Both take decorator parameters.
  • The middle below corresponds with the outer layer of __call__ above
  • we don’t have to use the nonlocal keyword because we can store values in self, like self.calltime
  • inner function is pretty similar across versions
def once_per_n(n): 
    def middle(f):
        calltime = 0 
        def inner(*args, **kwargs): 
            nonlocal calltime; 
            gap = time.time() - calltime 
            if gap < n: 
                msg = "You're calling this function too often. Try again in " + \
                          str(n - gap) + " seconds."
                raise Exception(msg)
            calltime = time.time()
            return f(*args, **kwargs)
        return inner  
    return middle

References / Further Reading Link to heading