Python Tutorial Creating a custom decorator

Posted in Uncategorized

In this Python tutorial, we will create a custom decorator named “print_start_time” which will wrap another function called “say_hello”.

The say_hello function is simply defined and called like this…

def say_hello(name):
    print("Hello " + name)

say_hello("World")

A custom decorator function will allow you to wrap this function like this …

@print_start_time
def say_hello(name):
    print("Hello " + name)

Which is equivalent to this …

def say_hello(name):
    print("Hello " + name)
say_hello = print_start_time(say_hello)

By doing this, it had alter the say_hello() function. Now when you call the say_hello() function, it will print the start time before printing “Hello”.

We define the custom decorator called “print_start_time” like this …

from datetime import datetime

def print_start_time(orig_func):
    def wrapper(*args, **kwargs):
        print("Start time: {}".format(datetime.now()))
        orig_func(*args, **kwargs)
    return wrapper

A decorator function is a function that takes another function as input parameter (in our case, orig_func) and returns a function (the newly modified function of the function passed in).

The returned function is typically called “wrapper” as we have above.

orig_func is the function the inner function that we want to wrap. But before calling it, it print the start time.

The wrapper function takes *args and **kwargs so that it can accept any parameters that the passed in function accepts. And it calls the passed function with those same arguments.

Using Python wraps

While the above code works, it’s introspection of decorated function is a bit broken. Ideally you should add the Python wraps decorator like this (after importing wraps from functools) …

from datetime import datetime
from functools import wraps

def print_start_time(orig_func):
    @wraps(orig_func)
    def wrapper(*args, **kwargs):
        print("Start time: {}".format(datetime.now()))
        orig_func(*args, **kwargs)
    return wrapper

This will preserve the __name__ and the __doc__ of your wrapped function. If you need to preserve more than that, read article here.

Writing decorator as class

Although less common, it is also possible to write our decorator as a class instead of a function…

class print_start_time():
    def __init__(self, orig_func):
        self.orig_func = orig_func

    def __call__(self, *args, **kwargs):
        print("Start time: {}".format(datetime.now()))
        return self.orig_func(*args, **kwargs)

---