A top-down approach to Python decorator design

This article started on my personal blog. For more Python, Django, and Vue development tutorials, please visit the Dream Chaser blog.

Decorators are an important programming practice in Python, but you can run into all sorts of difficulties when writing Python decorators if you don’t understand how they work and how to do it properly. I still remember that when I applied for the headliner Python development position in the school recruitment, I ended up with a design problem of Python decorator. I am very sorry. As MY programming skills improved and my understanding of Python decorators improved, I came up with a top-down approach to decorator design that made it easier to write all types of decorators without having to memorize decorator template code.

Python decorator principles

Here’s how Python decorators are normally written:

@decorator
def func(*args, **kwargs):
    do_something()
Copy the code

Inside the Python interpreter, calls to func are converted as follows:

>>> func(a, b, c='value')
# is equivalent to
>>> decorated_func = decorator(func)
>>> decorated_func(a, b, c='value')
Copy the code

As you can see, a decorator is a function (ora class) that takes the decorated function func as its only argument and returns a callable. The call to the decorated function func is actually a call to the returned Callable object.

Design decorators from the top down

As we can see from principle analysis, if we want to design a decorator that decorates the original function (or class) into a more powerful function (or class), all we need to do is write a function (or class) that returns the more powerful function (or class) we need.

Simple decorator

Simple decorator functions take no arguments, as described above. Suppose we want to design a decorator function whose function prints out the running time of the decorated function after the call. Let’s see how to design the decorator using the top-down approach.

The “top” is to not focus on implementation details, but to do the overall design and decomposition of the function call process. We’ll call the decorator timeThis and use it like this:

@timethis
def fun(*args, **kwargs):
    pass
Copy the code

Decomposing the call to decoratedfun:

>>> func(*args, **kwargs)
# is equivalent to
>>> decorated_func = timethis(func)
>>> decorated_func(a, b, c='value')
Copy the code

Our timeThis decorator should take the decorated function as its only argument and return a function object, which by convention is named Wrapper, so we can write the template code for the TimeThis decorator:

def timethis(func):
    def wrapper(*args, **kwargs):
        pass
    
    return wrapper
Copy the code

With the framework of the decorator in place, the next step is “down” to enrich the function logic.

Calling a decorated function is equivalent to calling the Wrapper function. To make the Wrapper call return the same result as the decorated function call, we can call the original function inside the Wrapper and return the result:

def timethis(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
    	return result
    return wrapper
Copy the code

We can vary the wrapper function logic as much as we want. We need to print the time of the func call, just before and after the func call:

import time

def timethis(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, end-start)
    	return result
    return wrapper
Copy the code

This completes a decorator that prints the time of a function call. Let’s see how it works:

@timethis
def fibonacci(n):
    Find the NTH term of the Fibonacci sequence
    a = b = 1
    while n > 2:
        a, b = b, a + b
        n -= 1
    return b

>>> fibonacci(10000)
fibonacci 0.004000663757324219. The result is too big to omitCopy the code

Basically, it looks fine, but since the function is decorated, the basic information about the decorated function becomes the information about the Wrapper function returned by the decorator:

>>> fibonacci.__name__
wrapper
>>> fibonacci.__doc__
None
Copy the code

Note that Fibonacci.__name__ is equivalent to timethis(Fibonacci).__name__, so return wrapper.

The fix is also simple: use a wraps decorator provided in the library to copy the wraps information to the Wrapper function:

from functools import wraps
import time

def timethis(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, end - start)
        return result

    return wrapper
Copy the code

At this point, a complete decorator with no parameters is written.

Decorator with parameters

The decorator design above is relatively simple and takes no parameters. We also often see decorators with parameters, which can be used as follows:

@logged('debug', name='example', message='message')
def fun(*args, **kwargs):
    pass
Copy the code

Decomposing the call to decoratedfun:

>>> func(a, b, c='value')
# is equivalent to
>>> decorator = logged('debug', name='example', message='message')
>>> decorated_func = decorator(func)
>>> decorated_func(a, b, c='value')
Copy the code

So, logged is a function that returns a decorator that decorates the func function, so the template code for logged should look something like this:

def logged(level, name=None, message=None):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            pass
        return wrapper
    return decorator
Copy the code

The wrapper is the final function to be called, and we can enrich and refine the decorator and Wrapper logic at will. Suppose our requirement is to print a line of log before the decorator function func is called as follows:

from functools import wraps

def logged(level, name=None, message=None):
    def decorator(func):
        logname = name if name else func.__module__
        logmsg = message if message else func.__name__

        @wraps(func)
        def wrapper(*args, **kwargs):
            print(level, logname, logmsg, sep=The '-')
            return func(*args, **kwargs)
        return wrapper
    return decorator
Copy the code

Multi-function decorator

Sometimes we’ll see two ways to use the same decorator, either as a simple decorator or by passing parameters. Such as:

@logged
def func(*args, **kwargs):
    pass

@logged(level='debug', name='example', message='message')
def fun(*args, **kwargs):
    pass
Copy the code

According to the previous analysis, a decorator with no parameters is defined differently than a decorator with parameters. A decorator with no arguments returns the decorated function. A decorator with arguments returns a decorator with no arguments, which then returns the decorated function. So how to unify? Let’s first analyze the invocation process of the two decorator methods.

Use @logged to decorate directly
>>> func(a, b, c='value')
# is equivalent to
>>> decorated_func = logged(func)
>>> decorated_func(a, b, c='value')

Logged (level='debug', name='example', message='message'
>>> func(a, b, c='value')
# is equivalent to
>>> decorator = logged(level='debug', name='example', message='message')
>>> decorated_func = decorator(func)
>>> decorated_func(a, b, c='value')
Copy the code

As you can see, the second decorator takes one step further than the first by calling the decorator function and returning a decorator that returns the same decorator as the one without arguments: it takes the function being decorated as its sole argument. The only difference is that the returned decorator has fixed partial parameters. Fixed partial parameters are exactly where partial functions are used, so we can define the following decorators:

from functools import wraps, partial

def logged(func=None, *, level='debug', name=None, message=None):
    if func is None:
        return partial(logged, level=level, name=name, message=message)

    logname = name if name else func.__module__
    logmsg = message if message else func.__name__

    @wraps(func)
    def wrapper(*args, **kwargs):
        print(level, logname, logmsg, sep=The '-')
        return func(*args, **kwargs)
    return wrapper
Copy the code

The key to the implementation is that if the decorator is used with arguments, the first argument func is None. We use partial to return a decorator with fixed arguments. This decorator, like a simple decorator with no arguments, takes the decorated function object as its only argument. It then returns the decorated function object.

Decoration class

Since the instantiation of a class is very similar to a function call, a decorator function can also be used to decorate a class, except that the first argument to the decorator function is no longer a function, but a class. Based on the top-down design approach, it is a snap to design a decorator function for a decorator class, which is not shown here.

practice

Finally, use today’s headline interview question as an exercise. Now it seems that this problem is just a simple decoration design needs, only blame their own learning skills are not fine, regret not to master the design method of decoration earlier.

Topic:

Design a retry decorator function. When a decorator call throws a specified exception, the function is recalled until the specified maximum number of calls is reached. The following is an example of a decorator:

@retry(times=10, traced_exceptions=ValueError, reraised_exception=CustomException)
def str2int(s):
    pass
Copy the code

Times is the maximum number of attempts for a function to be called again.

Traced_exceptions are monitored exceptions and can be None (the default), an exception class, or a list of exception classes. If None, all exceptions are monitored; If an exception class is specified, if the function call throws the specified exception, the function is re-called until the result is successfully returned or the maximum number of attempts is reached, at which point the original exception (reraised_exception has the value None) is re-thrown, or the exception specified by reraised_exception is thrown.

Reference code

Note that there is more than one implementation. Here is my version:

from functools import wraps

def retry(times, traced_exceptions=None, reraise_exception=None):
    def decorator(func):
        
        @wraps(func)
        def wrapper(*args, **kwargs):
            n = times
            trace_all = traced_exceptions is None
            trace_specified = traced_exceptions is not None
            while True:
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    traced = trace_specified and isinstance(e, traced_exceptions)
                    reach_limit = n == 0

                    if not (trace_all or traced) or reach_limit:
                        if reraise_exception is not None:
                            raise reraise_exception
                        raise
                    n -= 1
        return wrapper
    return decorator
Copy the code

conclusion

Sum up, decide and design decorator to divide the following steps

1. Decide how your decorator should be used, with or without arguments, or both.
2. Decomposing @ syntactic sugar into the actual invocation of the decorator.
3. Write the corresponding template code according to the invocation process of the decoration.
4. Write the logic for decorator functions and post-decorator functions as required.
5. Finished!

I share programming sentiment and learning materials of the public account, please pay attention to: programmer doughnuts