Decorators - Scientific Computing with Python

import numpy as np

We’ve seen Python functions and objects. Today, we’ll breifly cover decorators, which are functions that wrap functions, and come with sytactic sugar which lets you write things like

@time_wrapper
def my_function(x):
    ...
    
y = my_function(x)

> 1.1 seconds elapsed.

Again, we want to write a function that wraps another function. We can define this as

from time import monotonic

def time_wrapper(f):
    
    def inner_wrapper(*args, **kwargs):
        t0 = monotonic()
        ret = f(*args, **kwargs)
        t1 = monotonic()
        print("{} sec. elapsed".format(t1 - t0))
        return ret
    
    return inner_wrapper

Here’s how we might use our wrapper:

def generate_random(n):
    """
    return a random numpy vector
    """
    return np.random.rand(n)

generate_random = time_wrapper(generate_random)

x = generate_random(100000)
len(x)

one problem with this is that we had to go through the trouble of wrapping the function. Instead, we can just write

@time_wrapper
def generate_random2(n):
    """
    return a random numpy vector
    """
    return np.random.rand(n)

generate_random2(100000)

This is interpreted like what we did above, but is a bit easier to read.

Another problem we may encounter is that we now can’t access the docstring we wrote

help(generate_random2)

One way to solve this is to use the wraps decorator from the functools package.

from functools import wraps

def time_wrapper2(f):
    
    @wraps(f)
    def inner_wrapper(*args, **kwargs):
        t0 = monotonic()
        ret = f(*args, **kwargs)
        t1 = monotonic()
        print("{} sec. elapsed".format(t1 - t0))
        return ret
    
    return inner_wrapper

@time_wrapper2
def generate_random3(n):
    """
    return a random numpy vector
    """
    return np.random.rand(n)

generate_random3(100000)

now, we can access the metadata from generate_random3, such as the docstring

help(generate_random3)

As you might guess, the @wraps decorator copies the class metadata from f to the inner_wrapper function.