Function decorators

By Martin McBride, 2022-05-25

Tags: function decorator separation of concerns
Categories: python language intermediate python

Python decorators provide a simple way to modify or enhance existing functions. They are often used to add extra steps before or after a function in executed, which can be useful in many scenarios. There are other ways to use them too.

In this article we will review how decorators work, look at some simplified examples of how they can be used.

These are "toy" implementations, to show the basic concepts. In reality you would probably want to add extra functionality to make the decorators more flexible and to support more advanced language features. That is for a future article.

Logging

Here is a simple example function:

def show1():
    print("show1")

Suppose we wanted to add some logging functionality to this function. We will keep it simple, we just want to print a message to the console each time we enter the function, and each time we leave the function. We could do it like this:

def show1():
    print("Entering show1")
    print("show1")
    print("Leaving show1")

That works but it isn't ideal. We are cluttering our code with extra boilerplate code that has nothing to do with the main purpose of the function, which makes our code less readable. It would be better if our function just did its job, and secondary concerns such as logging were implemented somewhere else. This is called separation of concerns.

What if we could do this:

@log
def show1():
    print("show1")

@log
def show2():
    print("show2")

Here our function just does what it does. The magic @log statement tells Python to run some extra code before and after show1. We have also added a second function, show2 to show how the same logging boilerplate code can easily be added to any number of functions.

Implementing a decorator

Here is how we implement the decorator:

def log(function):

    def wrapped_function():
        print("Entering", function.__name__)
        function()
        print("Leaving", function.__name__)

    return wrapped_function

@log
def show1():
    print("show1")

@log
def show2():
    print("show2")

show1()
show2()

The log function is a special type of construct called a closure. We won't look at this in great detail, but here is what it does:

The log function accepts a parameter function.
Inside the log function we declare another function wrapped_function that implements the required logging functionality around whatever function we pass into log. wrapped_function is called an inner function of log.
The log function then returns wrapped_fnction.

The way the closure works is that when we call log(show1) it returns a brand new function (the inner function) that executes show1 with the logging code around it.

One final piece is the @ notation:

@log
def show1():
    print("show1")

This is a special notation that causes the Python interpreter to replace show1 with the result of log(show1). So now each time we call show1 we get the logging code automatically.

Here is the output:

Entering show1
show1
Leaving show1
Entering show2
show2
Leaving show2

Timing

Another use of decorators is to provide profiling information. We could easily modify the previous example to check the time at the start and end of each call, and store the timing data in some structure that can be accessed and analysed when the program ends. we won't show this here, it is very similar to the logging case.

Counting

Suppose we to count how many times each function is called. We can do this with a simple variation of the previous code:

def count(function):

    call_count = 0

    def wrapped_function():
        nonlocal call_count
        call_count += 1
        print("Call", call_count, "of", function.__name__)
        function()

    return wrapped_function

@count
def show1():
    print("show1")

@count
def show2():
    print("show2")

show1()
show2()
show1()
show2()

To count the number of calls, we need to store the current count between calls to the wrapped function. We do this by declaring a variable call_count in the main body of the count function. Within the wrapped function we have to access the variable by declaring to nonlocal, which tells Python to look in the containing function for that variable.

This relies on a very useful feature of closures - the inner function wrapped_function still has access to the variables in the outer function count even after the initial call to count has returned. wrapped_function maintains its own copy of the count function context.

Notice that @count is invoked twice, for show1 and show2. This means that the decorated versions of these functions each have their own context, including their own copy of call_count. So the code above will keep a separate count of calls to each function. Here is the output:

Call 1 of show1
show1
Call 1 of show2
show2
Call 2 of show1
show1
Call 2 of show2
show2

Checking parameter types

Here is an example function that requires a string argument:

def show_upper(s):
    print(s.upper())

If this function is called with a string, it will print the uppercase version of the string. If it is called with any other type of data, it will throw a cryptic error message that will probably mean nothing at all to the end user of the software.

We could add code to the function itself to check the type of the input parameter. That isn't a terrible thing to do, but it adds a bit of extra code that distracts form the main purpose of the function. Instead, we can add a decorator to check the parameter:

def accepts_string(function):

    def wrapped_function(s):
        if not isinstance(s, str):
            raise TypeError("String parameter required")
        return function(s)

    return wrapped_function

@accepts_string
def show_upper(s):
    print(s.upper())

show_upper("abc")
show_upper(1)

The decorator checks the parameter, and throws an exception if the parameter is not a string. You could enhance this by maybe throwing a custom exception, and perhaps logging a message too, but for the sake of simplicity we just throw a TypeError. The decorator adds this extra check without messing up the main code of the function. And of course you can reuse the same decorator on other functions, without duplicating code.

Checking parameter values

In the same vein, you can also use a decorator to check the value of a parameter, like this:

def check_value(function):

    def wrapped_function(s):
        if not s:
            raise ValueError("Non-empty string parameter required")
        return function(s)

    return wrapped_function

@check_value
def show_upper(s):
    print(s.upper())

show_upper("abc")
show_upper("")

This works in the same way to the previous example, but instead of checking the type of the argument it checks the value. Of course, in a real application you might want to check the type and the value. This can be combined in the same decorator.

Exception handling

Certain operations can cause runtime exceptions, and you will normally want to handle these in your code in some way. Typically, this is done using a try block to surround the code that might throw an exception. But of course that adds more boilerplate code to the function.

Once again we can consider useing a decorator to handle the error case, so that the function itself can stick to its main concern of executing the happy path, and ignoring the secondary concern of handling error cases. Here is the code:

def handle_exception(function):

    def wrapped_function(a, b):
        try:
            return function(a, b)
        except Exception:
            return None

    return wrapped_function

@handle_exception
def divide(a, b):
   return(a / b)

print(divide(1, 2))
print(divide(3, 0))

In this case our function performs a divide operation, which can potentially raise an exception if the divisor is zero.

The handle_exception decorator works by catching this exception and returning None. This means that our code doesn't need to bother with exception handling, but any code that calls divide will need to take account of the fact that a return value of None indicates an error.

There are other possibilities, for example we could log the exception, or rethrow the exception with additional information. In some cases it might even be preferable to ignore the exception altogether if it is non-critical, although that isn't usually a good idea.

Authentication

This example shows how we can avoid boilerplate code using decorators. In this first example, we will consider code that requires user authentication for certain operations:

def authenticate(function):

    def wrapped_function():
        if do_authenticate():
            function()

    return wrapped_function

@authenticate
def secure_function(s):
    print("Authorised user")

Here, secure_function is some function that should only be performed if the current user is authenticated. do_authenticate is a function that performs authentication and return true or false.

This implementation simply ignores the call is the user isn't known. You might wish to provide some kind of notification, or maybe throw an exception. This should be implemented within the wrapped_function.

Ignoring functions

These next few examples are probably things that you wouldn't want to do very often, but it is useful to know that they are possible.

This simple decorator causes a function to not be called:

def ignore(function):

    def wrapped_function():
        pass

    return wrapped_function

@ignore
def show():
    print("Executing show")

show() # Doesn't execute show()

Here wrapped_function doesn't call the original function. When we apply the ignore decorator to show, if we then call show, nothing will happen.

Switching functions

This decorator causes a different function to be called:

def switch(function):

    def wrapped_function():
        other_show()

    return wrapped_function

@switch
def show():
    print("Executing show")

def other_show():
    print("Executing other show")

show() # Actually executes other_show

Here wrapped_function always calls other_show regardless of which function is wrapped. When we apply the switch decorator to show, if we then call show, it will be other_show that runs.

Changing arguments and return values

We can use decorators to affect the function's arguments and return values.

This example implements byte subtraction. Byte values are limited to the range 0 to 255. In this example subtract_bytes does a simple subtraction. But the clamp decorator clamps the input values to the range 0 to 255 before calling the function, and clamps the result to the same range after the function returns. By clamping we mean that values less than 0 are set to 0 and values greater than 255 are set to 255.

def clamp(function):

    def inner(a, b):
        a = min(255, max(a, 0))
        b = min(255, max(b, 0))
        r = function(a, b)
        return min(255, max(r, 0))

    return inner


@clamp
def subtract_bytes(a, b):
    return a - b


print(subtract_bytes(500, 50))
print(subtract_bytes(120, 130))

In the first call, we pass in values of 500 and 50. But 500 is clamped to 255 so the result (255 - 50) is 205. In the second example we pass in 120 and 130, so the result (120 - 130) is -10, but that is clamped to 0 by the decorator.

Memoization

Suppose you have a function that takes a long time to execute, in a situation where it might be called any times with the same input values. In that situation, you might want to avoid calling the function more than once with the same input values.

Of course, this only works for pure functions (a pure function is a function that always returns the same value for a given set of inputs, and has no other side effects).

For example, the negate(x) function returns the negative of x. So if we call negate(1) it returns -1. But what if we call negate(1) again later? We don't need to repeat the calculation, our cade could simply remember the last time it was called with value 1, and return the same result.

This is called memoization, and here is an example of a decorator to do it:

def memoize(function):

    cache = dict()

    def inner(a):
        if a not in cache:
            cache[a] = function(a)
        return cache[a]

    return inner

@memoize
def negate(a):
    print("negating", a)
    return -a

print(negate(1))
print(negate(2))
print(negate(1))

The memoize decorator maintains a dictionary cache that holds previous results. The inner function checks if the result is known, and only calculates the result if it hasn't been calculated already. Here is the output from this code:

negating 1
-1
negating 2
-2
-1

Although we call negate(1) twice, and correctly print the result both times, the inner negate function is only called once (so the string "negating 1" is only printed once).

In reality, using memoization on such a simple function is probably not worthwhile, but with a complex and time-consuming functions this technique can lead to good efficiency improvements. Note also that this implementation uses the parameter value as a dictionary key, so our simple decorator only works with values that are suitable keys, such as numbers or strings. The functools module has a decorator called lru_cache that provides a full implementation.

Dynamic application of decorators

Remember that the @ notation is really just syntactic sugar - it makes the code nice to read, but you don't have to use it. You can apply the decorator manually, which allows you to create different versions of the same function. For example:

def log(function):

    def wrapped_function():
        print("Entering", function.__name__)
        function()
        print("Leaving", function.__name__)

    return wrapped_function

def show():
    print("show")

logged_show = log(show)

Remember that log(show) returns a function that executes show() surrounded by the logging code. So we can call logged_show() to execute show() with logging, but we can also call show() to execute the function directly with no logging.

Using decorators effectively

As we noted before, decorators are often used to implement separation of concerns, for example to separate logging or user authentication code from the main code.

This can be very useful but it can also be a double-edged sword. By their nature, decorators drag in hidden functions from distant corners of the code base, which has tha capacity to cause serious problems. It is usually best to limit this areas where you derive the most benefit, and generally use decorators that perform a conceptually simple function (such as logging or authenticating a user).

Basically, don't make decorators your go to solution for every problem. Use them sparingly, where the benefits definitely outweigh the downsides. But in the right scenario they are an extremely useful tool.