Automation Test Practice

Follow me on GitHub

Decorator in Python

In our daliy work when a project grows into a relative scale, very often we want to add additional features on an existing function. We can modify the function directly of course, however this violates Open-closed principle. Decorator is an effective solution to address such challenge–the name is pretty straightword to understand, when you decorator a window in your house, you donot change the functionality of the window. In Python, this can be implemented easily.

Essentially, Python decorator is a function or class. It adds additional functionalities for other method without changing the method code. We can use decorator for logging, authentication verification, and even to group test cases.

To start with a simple example as following, log_method is a decorator method. To use it on method hello_world(), we just use @ following with the decorator name.

def log_method(func):
    def wrapper():
        logging.info("%s is running" % func.__name__)
        return func()
    return wrapper

@log_method
def hello_world():
    print("hello from California")

Now I run method hello_world(), get following output:
python_1.jpg

As you can see, this decorator can be reused for many other methods.

This is still not robust enough. What if the business function hello_world has parameter(s)? What if the business function has default parameters? We need use ** *args ** and ** **kwargs **.

def log_method(func):

    def wrapper(*args, **kwargs):
        logging.info("%s is running" % func.__name__)
        return func(*args, **kwargs)
    return wrapper

Further more, decorator can have parameter(s). For example, I want to set different log lever when use the log_method decorator.

def log_method(level):
    def decorator(func):
        def wrapper(*args, **kwargs):
            if level == "info":
                logging.info("%s is running" % func.__name__)
            elif level == "warn":
                logging.warning("%s is running" % func.__name__)
            return func(*args, **kwargs)
        return wrapper
    return decorator

@log_method(level="info")
def hello_world():
    print("hello from California")

We basically add another layer of nested function to allow the decorator use parameter.

Python also has wraps from functools which can help to copy the meta information from the business method into the decorator.

@wraps(func)
def wrapper(*args, **kwargs):

Last but not least, decorator can be a class. We mainly rely on __call__ method to implement it.

class Decorator():
    def __init__(self, func):
        self._func = func

    def __call__(self):
        logging.info("this is in the decorator...")
        self._func()
        logging.info("decoration end here...")

@Decorator
def another_method():
    logging.info("this is in my real method...")

Back To Homepage