What is a decorator? Despite its name, a decorator doesn’t make python prettier. A decorator in python is a function that augments or modifies another function. Why is this useful? Well, firstly it saves you from having to write repeat code, because without a decorator you would probably have to include extra code in each function that does what the decorator does. Secondly, it allows you to do powerful things, such as add timers or error handlers to your code, easily and simply.
OK, I say simply, but some may find decorators a bit confusing, at least at first. So let’s take it step-by-step.
The Problem
Here’s a real-world use-case that a decorator can solve for us. Consider the following functions:
def droid(colour):
try:
if colour not in ['blue', 'red', 'gold']:
raise RuntimeError(f'No such droid colour "{colour}"')
return f'I am a {colour} droid'
except RuntimeError as ex:
print(str(ex))
def starfighter(colour):
try:
if colour not in ['red', 'blue', 'gold']:
raise RuntimeError(f'No such starfighter colour "{colour}"')
return f'Go faster {colour} stripe online'
except RuntimeError as ex:
print(str(ex))
The droid function attempts to define a droid for us, but raises exceptions if the droid colour is not valid. These exceptions are handled in the try:except block and the error message printed. Similarly for the starfighter function, exceptions are raised and handled if the starfighter colour is not valid.
We can run these functions with different arguments and print their return values:
new_droid_1 = droid('blue')
print(new_droid_1)
new_droid_2 = droid('purple')
print(new_droid_2)
new_starfighter_1 = starfighter('red')
print(new_starfighter_1)
new_starfighter_2 = starfighter('yellow')
print(new_starfighter_2)
The output we get is:
I am a blue droid No such droid colour "purple" None Go faster red stripe online No such starfighter colour "yellow" None
As you can see, the errors are handled and error messages printed. But we’ve put a lot of the same code in both functions for handling the errors. Could we use a decorator to handle the errors and reduce the code needed in each function? The answer, in case you hadn’t guessed, is yes!
The Goal
We are aiming to create a decorator called error_handler that will decorate the droid and starfighter functions above. Decorators are applied to functions by prefixing the decorator name with the ‘@’ symbol. The final syntax should look something like this:
@error_handler
def droid(colour):
...
@error_hander
def starfighter(colour):
...
Decorator Basics
Before we delve into creating a decorator, let’s get our head around some basics first. The syntax above is a shorthand way of writing what is really happening, which is this:
droid = error_handler(droid)
What this means is, the function droid is passed as an argument to the decorator error_handler and what is returned is an augmented version of droid that we still assign to the name droid. So when we call the decorated droid function, we end up calling the augmented version, thanks to the decorator.
In other words, because of the decorator, when we now call droid, we are actually calling error_handler(droid), i.e. making the call:
error_handler(droid)('some_colour')
Decorator Construction
Now we’ve got our heads round what is happening when we call a decorated function, what do we put in this magical decorator function?
Let’s start from the inside and work our way out. The core of our decorator is the part that handles the errors. It’s going to call our function, and wrap a try:except block around it. If no errors are caught, it will return the output from the function. If an error is caught, it will print the error message and return None. Here’s what it will look like:
def inner(*args, **kwargs):
try:
return func(*args, **kwargs)
except RuntimeError as ex:
print(str(ex))
The func function is the function that will be passed as an argument to the decorator. For our example, it will be the droid or starfighter function. The inner function will be an inner function of the decorator function. For more information on inner functions, see Real Python’s explanation on inner functions. The arguments passed to the inner function will be the arguments passed to the droid or starfighter functions, which is why we have used *args and **kwargs, so we can capture anything that is passed. Note that the inner function returns the output from func, which is needed if func returns output.
The decorator will wrap the inner function like so:
def error_handler(func):
def inner(*args, **kwargs):
...
return inner
Here, the decorator error_handler defines an inner function inner that does the actual augmentation to the original function that is passed as an argument to error_handler. The inner function is then returned when the error_handler decorator is called. From this we can see that these functions/calls are equivalent:
droid = error_handler(droid) error_handler(droid) = inner droid = inner
This means that when we call our decorated droid or starfighter functions, we are actually making this call:
inner('some_colour')
The inner function adds the error handling to our original function and prints or returns the result. The decorator in its entirety is:
def error_handler(func):
def inner(*args, **kwargs):
try:
return func(*args, **kwargs)
except RuntimeError as ex:
print(str(ex))
return inner
Now we have a decorator, we can remove the error handling from the droid and starfighter functions and apply the decorator to them. The final code is:
def error_handler(func):
def inner(*args, **kwargs):
try:
return func(*args, **kwargs)
except RuntimeError as ex:
print(str(ex))
return inner
@error_handler
def droid(colour):
if colour not in ['blue', 'red', 'gold']:
raise RuntimeError(f'No such droid colour "{colour}"')
return f'I am a {colour} droid'
@error_handler
def starfighter(colour):
if colour not in ['red', 'blue', 'gold']:
raise RuntimeError(f'No such starfighter colour "{colour}"')
return f'Go faster {colour} stripe online'
If we run the code in the same way as before:
new_droid_1 = droid('blue')
print(new_droid_1)
new_droid_2 = droid('purple')
print(new_droid_2)
new_starfighter_1 = starfighter('red')
print(new_starfighter_1)
new_starfighter_2 = starfighter('yellow')
print(new_starfighter_2)
The output we get is the same as before. Note that the output from the droid and starfighter functions was preserved by the decorator, as we are still able to print the output from the functions, even though they were decorated. This is because we returned the output from func in the inner function of the decorator.
I am a blue droid No such droid colour "purple" None Go faster red stripe online No such starfighter colour "yellow" None
And that’s it! Our droids and starfighters are decorated, and we can focus on their… erm, decoration without worrying about handling errors in each function. If you want to learn more about decorators, the Real Python Primer on Decorators has a good in-depth explanation. Otherwise, see you next time!
Recent Comments