Posts /

Python: Custom Function Decorators

Back to posts /
30 Aug 2016

Tutorial

You’ve probably seen function decorators in python before. They look something like this:

@my_decorator
def hello():
    print("Hello, world")

According to python’s documentation:

Decorators dynamically alter the functionality of a function, method, or class without having to directly use subclasses or change the source code of the function being decorated

An example of a library that makes heavy use of decorators (and which actually motivated me to create this post, in the first place) is Fabric. For example, Fabric has a decorator named runs_once, that can be used to make sure the decorated functions run only one time. For example:

@runs_once
def setup():
    # do setup, but only do it once!

That’s pretty cool, but what does that actually do? Well, decorators are simply functions that wrap other functions. That’s a slight oversimplification, but it will be our working definition for this article. In other words, these 2 examples are equivalent.

# 1
@my_decorater
def hello():
    print("Hello, world")
    
# 2
def hello():
    print("Hello, world")
hello = my_decorater(hello)

So, how can we implement our own decorators? Well, there are a few ways to accomplish this, but again, for this article, we will keep this as simple as possible. The method I prefer involves creating aDecoratorRegistry class, and using it to do the following:

Here is an example of the class I use. Add this to a file called DecoratorRegistry.py:

# DecoratorRegistry.py
from functools import wraps

class DecoratorRegistry:
    registry = {}

    @classmethod
    def registerDecorator(cls, name, func):
        if name not in cls.registry:
            cls.registry[name] = []
        cls.registry[name].append(func)

    @classmethod
    def getDecoratorFunctions(cls, name):
        if name in cls.registry:
            return cls.registry[name]
        return []

    @classmethod
    def callDecoratorFunctions(cls, name, *args, **kwds):
        decoratorFns = cls.getDecoratorFunctions(name)
        for fn in decoratorFns:
            fn(*args, **kwds) # call the decorator function

With this class defined, we can now define our own custom decorators pretty easily. For example, lets define a decorator named pre_deploy, that we can use to perform some intialization tasks before a code deployment. Add this function to the bottom of the DecoratorRegistry.py file (outside the class definition shown above):

# DecoratorRegistry.py
def pre_deploy(f):
    '''Defines a decorator named @pre_deploy'''
    @wraps(f)
    def wrapper(*args, **kwds):
        print("Running predeploy task: %s" % f(*args, **kwds))
    DecoratorRegistry.registerDecorator('pre_deploy', wrapper)
    return wrapper

Now that we’ve registered the pre_deploy decorator, we can add some functions in another file called deployment_tasks.py that use our new decorator.

# deployment_tasks.py
from DecoratorRegistry import pre_deploy

@pre_deploy
def step1():
    # perform some part of the setup
    return "step 1"

@pre_deploy
def step2():
    # perform some other part of the setup
    return "step 2"

Finally, lets add a file called deploy.py, which will handle our deployment logic, including the execution of ourpre_deploy tasks.

# deploy.py
from DecoratorRegistry import DecoratorRegistry
import deployment_tasks

print("Executing pre-deploy tasks")
DecoratorRegistry.callDecoratorFunctions('pre_deploy')

# add some deployment logic here
print("Deploy code")

print("Executing post-deploy tasks (to-do)")

Now, running this program from your shell will result in the following output:

$ python deploy.py

Executing pre-deploy tasks
Running predeploy task: step 1
Running predeploy task: step 2
Deploy code
Executing post-deploy tasks (to-do)

Next Steps