Tyler Poff


Decorators in Python!

06-02-2019

We've all had experience when we learn something new and the thought hit us: "why didn't I find out about this sooner?!" That simple thing that can make your life so much easier or solve a nagging problem you've been having. Something so obvious that had been staring you in your face your entire time. Something that, in retrospect, you probably should have learned about earlier, and you're a bit emberressed to admit you've only just found it. That thing that makes you go back and examine your past projects wondering how this new discovery could have made what you thought was a good design even better.

That was me a few weeks ago with decorators in Python.

What's a Decorator

If you've used Python for any length of time I'm sure you have used, or at least seen decorators, although you may not have thought much about them or even really understood what they were. I was first introduced to them when I started using Flask and making websites a few years ago.

The basic jist is a decorator is a design pattern that allows behavior to be added to an individual object without affecting it's original behavior. So, say you have a function that will sum up an array of integers and return a sum, like the function below

def sum_array(array):
    sum = 0 
    for number in array:
        sum += number
    return sum 

This is a beautiful function! (ok maybe not that beautiful....)

My point is, it does only 1 thing, it does that thing well. But say later on down the line your boss comes up to you and say "You know, I think that we need to add some time logging to that function, just to make sure it's not taking too long."

So you go back and rewrite it.

def sum_array(array):
    start_time = time.time()
    sum = 0 
    for number in array:
        sum += number
    end_time = time.time()
    elapsed_time = end_time-start_time 
    print("sum array took", elapsed_time, " seconds")
    return sum 

ok...still not too terrible, but it's starting to look a bit-"Just got a bug report, turns out our users are also passing in arrays that contain strings rather than numbers, we need to cast them"

def sum_array(array):
    for i in range(len(array)):
        array[i] = int(array[i])

    start_time = time.time()
    sum = 0 
    for number in array:
        sum += number
    end_time = time.time()
    elapsed_time = end_time-start_time 
    print("sum array took", elapsed_time, " seconds")
    return sum 

Well it could be wor-"Another requirement, now it needs to-"

And on, and one it goes, until you end up with a huge mess of a function. (Although let's be honest, alot of those things could have been done elsewhere in the code...). So, what can decorators do for us? they can help us add new functionality, the timing and the casting, without having to alter the original function. Better yet, if we design the decorators well enough we can use them elsewhere in the code for the same effect!

So let's create a simple decorator!

Creating a Decorator

So in Python, to implement a decorator you create a function, that takes in a function, and returns a function. Make sense? No? ok....

First things first, in Python you can pass and return functions as objects from functions. You can create pointers/variables of functions and then call them later in the code. Take the below for example.

def call_function(function_to_call):
    print("calling function.")
    function_to_call()
    print("function called!")
    
def function_1():
    print("Hello from function_1!")
def function_2():
    print("Hello from function_2")

call_function(function_2)
call_function(function_1)

The output from the above code is:

calling function.
Hello from function_2 
function called!
calling function.
Hello from function_1 
function called! 

So what's happening in the above code? We are passing a pointer to a function into call_function, call_function then calls the function the pointer is pointing towards. call_function doesn't know what function_to_call will be, or even what it will do. But whatever is passed in, it will try to call it as if it were a function. This is a powerful tool, and pretty cool in it's own right, and it's the basis for decorators.

In python a decorator will take a function as a parameter, if will then return a new function as a return variable. The idea is that this new function will wrap the old function and add any new functionality you may need. So a decorator takes in a function, wraps it inside a new function with any new functionality you need, and then return it. This new function will be called in place of the old one. So the old function gets extended using this wrapping method.

So let's re-write our sample from above and make call_function a decorator.

def first_decorator(function_to_call):
    def wrapper_function():
        print("calling function.")
        function_to_call()
        print("function called!")
    return wrapper_function 

@first_decorator
def function_1():
    print("Hello from function_1!")
@first_decorator
def function_2():
    print("Hello from function_2")
    
function_2()
function_1()

The output from the above code is:

calling function.
Hello from function_2 
function called!
calling function.
Hello from function_1 
function called! 

Which is the same as before. So let's break it down a bit to understand what it's doing.

We have a new function called first_decorator which takes the place of our call_function from before. This is our decorator function, just in case the name didn't give it away. It takes in a function as an argument, it has an inner function which is our wrapper, the wrapper will call the function_to_call along with our debug print statements. The first-decorator then returns this wrapper_function. Simple enough right?

Where the cool stuff comes in is when we get down to function_1 and function_2, you'll notice we now have a new line right above our function the "@first_decorator" line. What this line does is it tells python that, rather than just calling the function out right, we want to use the decorator to wrap this function. So everytime we call this function we are telling python to wrap it up and then call the wrapped function. Python will take the hint and make sure that it goes through our new wrapper whenever we want to call the function_1 or function_2, or really any function that has that @first_decorator tag before it.

This is powerful syntax that python gives us and allows us to add decorators easily.

So if we took away the python syntax we might have

new_function = first_decorator(function_1)
new_function() 

Every time we call function_1(), the @first_decorator is just telling python we want to do that, and makes our code cleaner and more readable.

Decorators and function arguments

You'll notice that our functions and decorator wrappers haven't been taking any arguments, but never fear, it is very easy to do just that.

Going back to our first example, let's re-write our sum_array function and add a decorator to time how long it took for the function to run.

import time
def timing_decorator(function_to_call):
    def wrapper_function(array_to_sum):
        start_time = time.time()
        sum = function_to_call(array_to_sum)
        end_time = time.time()
        elapsed_time = end_time-start_time 
        print("sum array took", elapsed_time, " seconds")
        return sum 
    return wrapper_function
        
@timing_decorator
def sum_array(array):
    sum = 0 
    for number in array:
        sum += number
    return sum 
        
array_to_sum = [1, 2, 3]*50000
sum = sum_array(array_to_sum)
print("sum is ",sum)

So in the above example, our decorator still takes in a single function as an argument, that remains the same. But our INNER function, the wrapper_function, now takes a single argument, which will be our array_to_sum. It then passes this argument onto it's function that it's wrapping.

And it's that simple, just by defining the inner function of our decorator to take an argument is all that's needed. HOWEVER....there is a cooler way to do that.

So say we want to have our timing decorator be able to be added to ANY function, regardless of how many arguments it has? Doesn't take any arguments? our decorator can handle it. Takes over 100? I...would question why you would write a function like that but sure.

How can we do that? by using the *args and **kwargs arguments in our decorator.

*args and **kwargs

So the *args and **kwargs are idioms in python for packing and unpacking arguments. You use these variables when you aren't sure how many variables will be passed into a function. This is useful for our wrapping functions since this means we can make decorators that wrap any type of function.

So how do they store the variables? Simply put the *args variable stores all non-keyworded variables and stores them into a tuple, while **kwargs stores all keyworded variables and stores them in a dict structures.

Take a look at the below example that uses *args and **kwargs.

# this function will just print the various values for *args and **kargs
def print_args_kargs(*args, **kargs):
    print("args values: ",args)
    print("kargs values: ",kargs)
    
#               |     non-keyworded variables        |     key-worded variables
print_args_kargs(1, 200, 5000, "non-keyworded string", test2="passing in named variables are cool!", secret_number=4)

The above sample will output.

args values: (1, 200, 5000, "non-keyworded string")
kargs values: {"test2":"passing in named variables are cool!", "secret_number":4}

So as you can see, any variables that don't have a name specified goes to *args, while named variables go into **kargs. But a major point I want to point out is that you can pass in *args and **kwargs into a function call. This is very useful if you are writting wrapper functions (like our decorators) or other functions where you aren't quite sure what variables will be passed in.

Cool, another neat feature we can add to our python toolkit, let's see how we can integrate this into our decoraters, we'll re-write our timing decorator from before using *args and **kwargs.

import time
def timing_decorator(function_to_call):
    def wrapper_function(*args, **kwargs):
        start_time = time.time()
        sum = function_to_call(*args, **kwargs)
        end_time = time.time()
        elapsed_time = end_time-start_time 
        print("function ",function_to_call.__name__," took ", elapsed_time, " seconds to run")
        return sum 
    return wrapper_function
        
@timing_decorator
def sum_array(array):
    sum = 0 
    for number in array:
        sum += number
    return sum 
@timing_decorator
def print_larger_number(first, second):
    if first > second:
        print(first)
    else:
        print(second)

array_to_sum = [1, 2, 3]*50000
sum = sum_array(array_to_sum)
print("sum is ",sum)
print()
print_larger_number(50, 25)

There we have it. By just changing the inner wrapper_function in our decorator to accept *args and **kwargs as arguments we can use it to wrap ANY function in python. This gives use a timing decorator that we can use anywhere in our code and for any function or class method.

Decorator Arguments

So allready I hope you can see how powerful and useful decorators can be. They can be used as timers, or as loggers, or to validate variables, there are numerous possibilities for these useful python tools. Another thing I want to show you about decorators is how to pass in arguments to the decorator itself.

So say for example you want to write a decorator to print a warning if a function returns a value of False, but the message needs to be different for each function. We could write our function so that it returns 2 values, a False and an error message and then check in our calling code if it's false and then print it out....OR, we could write a decorator that takes in the failure/warning message as a variable, and does all the checking in there. I like that design better.

The point of this example is to show you that you can pass arguments and variables into decorators. This can be very useful especially if the decorator needs them to determine what action to take (maybe take in a "failure" function to call or like in our example a failure message). In order to do THAT, we're going to add ANOTHER layer to our decorator function.

So, let's see what it looks like.

''' this is our decorator, it has an inner function which will accept a function to call, it also has a inner function which is our wrapper.
 the only difference is now we have this upper level function which takes in an argument, in this case our warning to print
 Other than that it acts the same way, python will still use it as a wrapper for whatever function we decorate it with.'''
def print_warning_on_false_decorator(warning):
    def actual_decorator(function_to_call):
        def wrapper_function(*args, **kwargs):
            # if our function returns false, print our error message that we got from the highest level of our function. 
            result = function_to_call(*args, **kwargs)
            if result == False:
                print(warning)
            return result 
        return wrapper_function
    return actual_decorator
@print_warning_on_false_decorator("this won't ever get called")
def return_true():
    return True
@print_warning_on_false_decorator("on no, the return false function returned false? what a surprise.")
def return_false():
    return False

# now call our 2 functions, 
return_true()
return_false()

It may look confusing, but it's very similar to what we've been doing. It just adds another layer which holds our argument. In the above example each function, return_true and return_false has their own unique warning, if they return false our decorator will print their unique warnings which we define in the decorator line. We pass this warning in just like an argument to a function. This is extremely useful for customizing a decorators response or action for whatever function it wraps.

Decorator Classes

Last thing I want to mention is that you can write decorators as classes. Up until now we've just written them as functions, but this isn't the only option python gives us. If we have a particularly complex decorator, say with states we want to keep track of for whatever reason, we can write it as a class and it can hold those variables for us.

Here's a quick example:

class Class_As_Decorator():
    def __init__(self, function_to_call):
        self.counter = 0
        self.function_to_call = function_to_call
        print("decorator as class initialized.")
    def __call__(self, *args, **kwargs):
        self.counter += 1
        print(self.function_to_call.__name__, " has been called ",self.counter," times")
        return self.function_to_call(*args, **kwargs)
        
@Class_As_Decorator
def sum_array(array):
    sum = 0 
    for number in array:
        sum += number
    return sum 

array_to_sum = [1, 2, 3]*50000
sum = sum_array(array_to_sum)
sum = sum_array(array_to_sum)
sum = sum_array(array_to_sum)
sum = sum_array(array_to_sum)
sum = sum_array(array_to_sum)
print("sum is ",sum)

So in this example we write a class decorator that just keeps track of how many times the decorated function has been called. But this shows the basic outline for how to define a class decorator and how it can keep track of it's own state and variables. Pretty cool right?

In Conclusion

So that's Python decorators, there's alot more you can explore about them but this should serve as a good introduction. They are incredibly useful and I find myself using them more and more to help modularize my code and keep functions clean. Hopefully this was of some use to you and I hope you learned something from this blog post! Happy Coding!

Page Links