How to write DRY code in Python using decorator functions!

How to write DRY code in Python using decorator functions!
Photo by Mike Erskine on Unsplash

Update: You can now view this post in video format here:

Writing good, scalable and maintainable software is an art and a must have skill for any aspiring software engineer. There are many principles out there including SOLID or GRASP.
The objective of this post is to focus on a much simpler, single principle called DRY which stands for don’t repeat yourself. Copy-pasting code in multiple locations in the same project introduces unnecessary complexity and makes the code prone to bugs. Debugging becomes harder and so does testing.

One brilliant way to embody this principle in Python is to make use of decorator functions. Let’s dive in!

The problem of repeated code

Let us create two functions simpleArith and complexArith as follows:


def simpleArith(num1, num2, operator):
    if operator == '+':
        return {"result": num1 + num2}
    elif operator == '-':
        return {"result": num1 - num2}
    elif operator == '*':
        return{"result": num1 * num2}
    elif operator == '/':
        return {"result": num1 / num2}


def complexArith(num1, num2, num3, operator):
    if operator == '+':
        return {"result": num1 + num2 + num3}
    elif operator == '-':
        return {"result": num1 - num2 - num3}
    elif operator == '*':
        return {"result": num1 * num2 * num3}
    elif operator == '/':
        return {"result": num1 / num2 / num3}

The functions are pretty straightforward. You can go ahead and test them by adding the following line to the end of your .py file:

print(simpleArith(1, 2, '+'))

All well and good. But if you try to pass in an operator that is not supported the program crashes with an error:

print(simpleArith('1', 2, '+'))
>> (torchenv) PS C:\Users\hasee\Documents\blog\decorator> py .\fun.py
Traceback (most recent call last):
  File "C:\Users\hasee\Documents\blog\decorator\fun.py", line 49, in <module>
    print(simpleArith('1', 2, '+'))
  File "C:\Users\hasee\Documents\blog\decorator\fun.py", line 30, in simpleArith
    return {"result": num1 + num2}
TypeError: can only concatenate str (not "int") to str

The fix is fairly easy in this case but the point is that it needs to be copy-pasted to both simpleArith and complexArith.
However, if we do this we would be repeating our code. This is bad practice.

Like so:

def simpleArith(num1, num2, operator):
    
    try:
        int(num1)
    except ValueError:
        return {'message': 'Invalid input'}
    
    try:
        int(num2)
    except ValueError:
        return {'message': 'Invalid input'}

    if operator == '+':
        return {"result": num1 + num2}
    elif operator == '-':
        return {"result": num1 - num2}
    elif operator == '*':
        return{"result": num1 * num2}
    elif operator == '/':
        return {"result": num1 / num2}

def complexArith(num1, num2, num3, operator):

    try:
        int(num1)
    except ValueError:
        return {'message': 'Invalid input'}
    
    try:
        int(num2)
    except ValueError:
        return {'message': 'Invalid input'}

    try:
        int(num3)
    except ValueError:
        return {'message': 'Invalid input'}

    if operator == '+':
        return {"result": num1 + num2 + num3}
    elif operator == '-':
        return {"result": num1 - num2 - num3}
    elif operator == '*':
        return {"result": num1 * num2 * num3}
    elif operator == '/':
        return {"result": num1 / num2 / num3}

Notice how much repetition of code there is. The try-catch block is repeated 5 times!

We can do better 😊

Before we jump to the solution, we need to review one important concept in Python. Namely that of higher-order functions.

Passing functions to functions

Before we move onto decorator functions, we need to simplify one concept first. In Python it is possible to pass a function as an argument to another function. This is because functions in Python are actually higher-order functions.

Consider the following code:

def fun(arg):
    arg()

def fun2():
    print('fun2')

fun(fun2)

What do you think the output of this would be? That’s right. The arg variable is a reference to the fun2 function and it can be called as you normally would call a function.
We see the string ‘fun2 printed:

passing functions in functions

Decorator functions

A decorator function is nothing more than a function that takes a function as argument.
The basic structure looks like the following:

def decorator(fun):
    def wrapper(*args, **kwargs):
      
        return fun(*args, **kwargs)
    return wrapper

As you can see, a function (fun) is passed as an argument to the decorator function. The wrapper function is then invoked which has *args and **kwargs as arguments. The *args variable is a tuple containing the arguments passed to the original function fun while the **kwargs is a dictionary containing the keyword arguments passed to the fun function.

This will make more sense in a bit.

If we use the @ notation before a function definition, the decorator function will wrap around the function being defined. We can remove the error checking logic in simpleArith and add @decorator like so:

@decorator
def simpleArith(num1, num2, operator):
 
    if operator == '+':
        return {"result": num1 + num2}
    elif operator == '-':
        return {"result": num1 - num2}
    elif operator == '*':
        return{"result": num1 * num2}
    elif operator == '/':
        return {"result": num1 / num2}

Whenever simpleArith is invoked, decorator will be invoked first with simpleArith as the argument. The arguments num1, num2, operator2 will appear as a tuple in the *args variable.

We can now put the error checking logic into the decorator function like so:

def decorator(fun):
    def wrapper(*args, **kwargs):
      
        nums = args[:-1]
        for arg in nums:
            try:
                int(arg)
            except ValueError:
                return {'message': 'Invalid input'}

        tu=tuple(map(int, nums))
        tu = tu + (args[-1],)
        return fun(*tu, **kwargs)

    return wrapper

We are removing the last argument (which is the operator) and storing the result in nums. Then we are iterating over this and trying to do a cast to integer. If this fails we know that the input is bad and we can catch the error and inform the user that the input is invalid.

If the input is ok, we cast the numbers from string to int just to make sure the program doesn’t crash if an input like ‘1’ is passed. We use the map function:

tu=tuple(map(int, nums))

However, we also need to append the operator which is contained in args[-1] to this like so:

tu = tu + (args[-1],)

When this is done we return by calling fun with our modified arguments. In this case fun is the original function being decorated which is simpleArith. So we are calling simpleArith passing the casted arguments to it.

return fun(*tu, **kwargs)

That’s it. It is very simple. Now here comes the magic. Suppose we want the same logic to apply to the complexArith function. What do we do? We simply wrap it with @decorator like so:

@decorator
def complexArith(num1, num2, num3, operator):

    if operator == '+':
        return {"result": num1 + num2 + num3}
    elif operator == '-':
        return {"result": num1 - num2 - num3}
    elif operator == '*':
        return {"result": num1 * num2 * num3}
    elif operator == '/':
        return {"result": num1 / num2 / num3}

Now we can call both simpleArith and complexArith with bad inputs and the decorator would be invoked first in both cases, taking care of the input!

Code is available here: fun.py (github.com)

Recap

Phew that was a lot!

To summarize, a decorator function can be thought of as a wrapper that wraps around a function. This means, if that function is called, the decorator gets called first. The function is passed to the decorator as an argument. Thus, arguments to that function can also be accessed from within the decorator via *args and **kwargs.

Decorator functions are useful to make DRY code which avoids code repetition.
They are used in cases where the same logic needs to be applied to multiple functions. Because the *args and **kwargs types are used, we do not restrict ourselves to functions that have only one, two or three arguments. We can decorate a function with any number of arguments. This makes decorator functions so versatile to use!

I post tidbits on Twitter related to Python. Consider following me here: Haseeb Kamal (@mhaseebkamal) / Twitter
You can also find me on LinkedIn here: Muhammad Haseeb Kamal | LinkedIn

You can find more code here.

Have a great day!