Wrap a function using Python Decorators

What is Decorator in Python and what problem decorators will try to solve.

Let me explain with an example.

Here we have 2 functions , calculate square and calculate cube.

def calc_square(numbers):
    result = []
    for number in numbers:
        result.append(number*number)
    return result


def calc_cube(numbers):
    result = []
    for number in numbers:
        result.append(number*number*number)
    return result

array = range(1,100000)
out_square = calc_square(array)
out_cube = calc_cube(array)

Here, calc_square function takes array of numbers as an input and iterating through the array ,calculating the square of a number and putting that into a result.

calc_cube does same thing, but instead of square ,it does cube. Here we are calling those function for a range of 1 to 100000.

Now often you have a need of measuring the performance of a function. Performance, I mean how much time does every function take to execute. In order to measure the timing, you have to use the time module. I will take the start time . Once you are done ,you will take end time. The code snippet is shown below.

import time

def calc_square(numbers):
    start = time.time()
    result = []
    for number in numbers:
        result.append(number*number)
    end = time.time()
    print(" calc_square took " + str((end-start)*1000) + " mil sec ")
    return result

We do the same thing for cube function as well, because we want to measure the performance of both of these functions .

def calc_cube(numbers):
    start = time.time()
    result = []
    for number in numbers:
        result.append(number*number*number)
    end = time.time()
    print(" calc_cube took " + str((end-start)*1000) + " mil sec ")
    return result

The output of the above code will be

calc_square took 27.922630310058594 mil sec 
calc_cube took 41.88823699951172 mil sec

Now the problem with this code is that , lets say you have a complex software project and you have written more than 100 functions. In order to measure the performance of all those functions , you have to write the start time and end time , exactly the same line of code in every function .

Problem No.1 is The start time and end time is getting repeated in every function that you want to measure the performance.

Problem no.2 is , there is a logic in the function(calculating square and cube) and it is combined with the timing logic .It makes code less readable.

Now there has to be a better way of doing this and better way is basically a Decorator.

Decorator allows you to wrap your function in another function .
Lets remove timing logic or timing code from above functions. And I want to have a function which has just the logic that function is supposed to do.

In order to do decorator , you need to define the wrapper function and lets call it as time_it(func) and that wrapper function will take function (func) as an argument.

Now functions are 1st class objects in Python means:

• You can pass functions as an argument to function.

• Also, you can return function as a return value from another function.

So in time_it(func) : we will define another function called wrapper (Python allows us to write nested function i. e we can write one function inside another function) and what wrapper function is doing is it is taking the positional arguments *args and keyword arguments which is **kwargs

And then it will start the timer and then it will call the function that was passed as an argument (So I am going to call function here) with argument and keyword argument

Then I will measure the end time.

The print(func.__name__ + “ took ” + str((end-start)**1000 + “ mil sec ”)

then return result

Then return wrapper function. i.e return wrapper

 __name__

will return the name of that function. Here we are returning a function (wrapper) from another function time_it , that’s why this function wrapper is called 1st class object . You can treat it as normal variable, you can return it and pass it as a function argument and so on.

def time_it(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args,**kwargs)
        end = time.time()
        print(func.__name__ + " took " + str((end-start)*1000) + " mil sec ")
        return result
    return wrapper

Also we need to decorate the calc_square and calc_cube function i.e @time_it which is shown below

@time_it
def calc_square(numbers):
    result = []
    for number in numbers:
        result.append(number*number)
    return result

@time_it
def calc_cube(numbers):
    result = []
    for number in numbers:
        result.append(number*number*number)
    return result

So, any function that you want to measure the performance of now, once you have defined this time_it function , you can put the @time_it tag the beginning and it is going to measure the performance . It makes both the function calc_square and calc_cube more readable and all your timing code is restricted into one function.

Final thoughts :

Decorator acts as a wrapper to the original function.

waughkes