Decorator: Advance

In Decorator: Basics, we have seen the basics of decorators.

import functools

def decorator(func):
    @functools.wraps(func)
    def wrapper_decorator(*args, **kwargs):
        # Do something before
        value = func(*args, **kwargs)
        # Do something after
        return value
    return wrapper_decorator

The key is to remember the following two expressions are equivalent (the latter uses the syntactic sugar):

def function():
    pass

function = decorator(function)
@decorator
def function():
    pass

Here we will explore more advanced features

We will resue some custom decorators defined before:

import functools

def debug(func):
    """Print the function signature and return value"""
    @functools.wraps(func)
    def wrapper_debug(*args, **kwargs):
        args_repr = [repr(arg) for arg in args]
        kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
        signature = ", ".join(args_repr + kwargs_repr)
        print(f"Calling {func.__name__}({signature})")
        value = func(*args, **kwargs)
        print(f"{func.__name__!r} returned {value!r}")
        return value

    return wrapper_debug


def timer(func):
    """Print the runtime of the decorated function"""
    @functools.wraps(func)
    def wrapper_timer(*args, **kwargs):
        start_time = time.perf_counter()
        value = func(*args, **kwargs)
        end_time = time.perf_counter()
        runtime = end_time - start_time
        print(f"Finished {func.__name__!r}: in {runtime:.4f} secs.")
        return value

    return wrapper_timer


def do_twice(func):
    @functools.wraps(func)
    def wrapper_do_twice(*args, **kwargs):
        func(*args, **kwargs)
        return func(*args, **kwargs)
    return wrapper_do_twice

Decorating Classes

There are two different ways to use decorators on classes

Decorate the methods of a class

Some commonly used decorators (that are even built-ins in Pythons)

  • @classmethod, @staticmethod

    Define methods inside a class namespace that are not connected to a particular instance of that class

  • @property

    Customize getters and setters for class attribute

Example
class Circle:

    def __init__(self, radius):
        self._radius = radius

    # radius is a mutable property: it can be set to a different value.
    @property
    def radius(self):
        """Get value of radius"""
        return self._radius

    @radius.setter
    def radius(self, value):
        """Set radius, raise error if negative"""
        if value >= 0:
            self._radius = value
        else:
            raise ValueError("Radius must be positive")

    # area is an immutable property 
    # since preperties without setter methods can not be changed
    @property
    def area(self):
        """Calculate area inside circle"""
        return self.pi() * self.radius ** 2

    # regular method
    def cylinder_volume(self, height):
        """Calculate volume of cylinder with circle as base"""
        return self.area * height

    # Class method, which is not bound to one particular instance of Circle.
    # Class methods are often used as factory methods that can create specific instances of the class.
    @classmethod
    def unit_circle(cls):
        """Factory method creating a circle with radius 1"""
        return cls(1)

    # Static method.
    # It’s not really dependent on the Circle class, except that it is part of its namespace. 
    # Static methods can be called on either an instance or the class.
    @staticmethod
    def pi():
        """Value of Ο€, could use math.pi instead though"""
        return 3.14155926535
>>> c = Circle(5)
>>> c.radius
5

>>> c.area
78.5398163375

>>> c.radius = 2
>>> c.area
12.566370614

>>> c.area = 100
AttributeError: can't set attribute
>>> c.cylinder_volume(height=4)
50.265482456

>>> c.radius = -1
ValueError: Radius must be positive

>>> c = Circle.unit_circle()
>>> c.radius
1

>>> c.pi()
3.1415926535

>>> Circle.pi()
3.1415926535

We can also apply custom decoraters to decorate methods.

Example
class TimeWaster:
    @debug
    def __init__(self, max_num):
        self.max_num = max_num

    @timer
    def waste_time(self, num_times):
        for _ in range(num_times):
            sum([i**2 for i in range(self.max_num)])
>>> tw = TimeWaster(1000)
Calling __init__(<time_waster.TimeWaster object at 0x7efccce03908>, 1000)
'__init__' returned None

>>> tw.waste_time(999)
Finished 'waste_time' in 0.3376 secs

Decorate the whole class

A common use of class decorators is to be a simpler alternative to some use-cases of metaclasses. For example, the new dataclasses module in Python 3.7

from dataclasses import dataclass

@dataclass
class PlayingCard:
    rank: str
    suit: str

The meaning of the syntax is similar to the function decorators: In the example above, you could have done the decoration by writing PlayingCard = dataclass(PlayingCard).

Writing a class decorator is very similar to writing a function decorator. The only difference is that the decorator will receive a class and not a function as an argument.

Example
@timer # a shorthand for TimeWaster = timer(TimeWaster)
class TimeWaster:
    def __init__(self, max_num):
        self.max_num = max_num

    def waste_time(self, num_times):
        for _ in range(num_times):
            sum([i**2 for i in range(self.max_num)])

Decorating a class does NOT decorate its methods. Here, @timer only measures the time it takes to instantiate the class:

>>> tw = TimeWaster(1000)
Finished 'TimeWaster' in 0.0000 secs

>>> tw.waste_time(999)
>>>

Nesting Decorators

You can apply more than one decorator to a function. This accumulates their effects and it’s what makes decorators so helpful as reusable building blocks.

The order of application of decorators is from bottom to top. It is sometimes called decorator stacking: You start building the stack at the bottom and then keep adding new blocks on top to work your way upwards.

Example
@debug
@do_twice
def greet(name):
    print(f"Hello {name}")

The decorators will be executed in the order they are listed. In other words, @debug calls @do_twice, which calls greet(), or debug(do_twice(greet())).

>>> greet("Eva")
Calling greet('Eva')
Hello Eva
Hello Eva
'greet' returned None

Change the order of @debug and @do_twice:

from decorators import debug, do_twice

@do_twice
@debug
def greet(name):
    print(f"Hello {name}")

In this case, @do_twice will be applied to @debug as well (do_twice(debug(greet()))):

>>> greet("Eva")
Calling greet('Eva')
Hello Eva
'greet' returned None
Calling greet('Eva')
Hello Eva
'greet' returned None
Example

We define two decorators which wrap the output string of the decorated function in HTML tags.

def strong(func):
    def wrapper():
        return f"<strong> {func()} </strong>"
    return wrapper


def emphasis(func):
    def wrapper():
        return f"<em> {func} </em>"
    return wrapper

Apply them to the greet() function at the sme time:

@strong
@emphasis
def greet():
    return "Hello!"
>>> greet() 
'<strong><em>Hello!</em></strong>'

Decorators With Arguments

Syntax for decorators with parameters:

@decorator(params)
def func_name():
    ''' Function implementation'''

The above code is equivalent to

def func_name():
    ''' Function implementation'''

func_name = (decorator(params))(func_name)

To create a decorator that accepts arguments, you need to create a “meta-decorator” function that

  • takes arguments and
  • returns a regular decorator (which in turns returns a function)

Example: Extend @do_twice to a @repeat(num_times) decorator, where the number of times to execute the decorated function could then be given as an argument.

import functools

def repeat(num_times):
    # By passing num_times, a closure is created, 
    # where the value of num_times is stored until it will be used later by wrapper_repeat()
    
    ### Just a common regular decorator ###
    def decorator_repeat(func):
        @functools.wraps(func)
        def wrapper_repeat(*args, **kwargs):
            for _ in range(num_times):
                value = func(*args, **kwargs)
            return value
        return wrapper_repeat
    #######################################
    
    return decorator_repeat
	
@repeat(num_times=4)
def greet(name):
    print(f"Hello {name}")
>>> greet("World")
Hello World
Hello World
Hello World
Hello World

More flexible: optionally take arguments

We can also define decorator that can be used both with and without arguments.

import functools

def name(
    original_func=None, 
    *, # enforce that following parameters are keyword-only
    kw1=val1,
    kw2=val2,
    ...
):
    def decorator_name(function):
        @functools.wraps(function)
        def wrapper_function(*args, **kwargs):
            ...
            
        return wrapper_function
    
    if original_func:
        return decorator_name(original_func)
    
    return decorator_name
    
Explanation
  • When the name is called with no optional arguments

    @name
    def function():
        ...
    

    The function is passed as the first argument and the decorated function will be returned:

    function = name(original_func=function)
    

    As name(original_func=function) returns decorator_name(function), the code above is equivalent to

    function = decorator_name(function)
    

    which is the same as a regular decorator.

  • When the decorator is called with one or more optional arguments

    @name(kw1="some value")
    def function():
        ...
    

    The name is called with original_func=None and kw1="some value":

    function = (name(original_func=None, * kw1="some value"))(function)
    

    As name(original_func=None, * kw1="some value") returns decorator_name, the code above is equivalent to

    function = decorator_name(function)
    

    , as expected.

Example: Apply this boilerplate on the @repeat decorator

def repeat(original_func=None, *, num_times=2):
    def decorator_repeat(func):
        @functools.wraps(func):
            def wrapper_repeat(*args, **kwargs):
                for _ in range(num_times):
                    value = func(*args, **kwargs)
                return value
            return wrapper_repeat
        
    if original_func:
        return decorator_repeat(original_func)
    
    return decorator_repeat
@repeat
def say_whee():
    print("Whee!")

@repeat(num_times=3)
def greet(name):
    print(f"Hello {name}")
>>> say_whee()
Whee!
Whee!

>>> greet("Penny")
Hello Penny
Hello Penny
Hello Penny

Stateful Decorators

Sometimes it’s useful to have a decorator that can keep track of state. We can use the function attribute of the wrapper function to store the state.

Example: create a decorator that counts the number of times a function is called.

import functools

def count_calls(func):
    @functools.wraps(func)
    def wrapper_count_calls(*args, **kwargs):
        wrapper_count_calls.num_calls += 1 # update the state
        print(f"Call {wrapper_count_calls.num_calls} of {func.__name__!r}")
        return func(*args, **kwargs)
    wrapper_count_calls.num_calls = 0 # store the initialized state using the function attribute
    return wrapper_count_calls

@count_calls
def say_whee():
    print("Whee!")
>>> say_whee()
Call 1 of 'say_whee'
Whee!

>>> say_whee()
Call 2 of 'say_whee'
Whee!

>>> say_whee.num_calls
2

Classes as Decorators

The typical way to maintain state is by using classes.

A typical implementation of a decorator class needs to implement

  • .__init__()
    • stores a reference to the function
    • do any other necessary initialization
  • .__call__(): the class instance needs to be callable so that it can stand in for the decorated function

Example: implement count_calls in the previous section

import functools

class CountCalls:
    def __init__(self, func):
        functools.update_wrapper(self, func)
        self.func = func
        self.num_calls = 0

    def __call__(self, *args, **kwargs):
        self.num_calls += 1
        print(f"Call {self.num_calls} of {self.func.__name__!r}")
        return self.func(*args, **kwargs)

@CountCalls
def say_whee():
    print("Whee!")
>>> say_whee()
Call 1 of 'say_whee'
Whee!

>>> say_whee()
Call 2 of 'say_whee'
Whee!

>>> say_whee.num_calls
2

Real World Examples / Use Cases

Slowing Down Code

import functools
import time

def slow_down(_func=None, *, rate=1):
    def decorator_slow_down(func):
        @functools.wraps(func)
        def wrapper_slow_down(*args, **kwargs):
            time.sleep(rate)
            print(f"Sleep for {rate} second.")
            value = func(*args, **kwargs)
            return value
        return wrapper_slow_down
    if _func:
        return decorator_slow_down(_func)
    
    return decorator_slow_down
@slow_down(rate=2)
def countdown(from_number):
    if from_number < 1:
        print("Lift off!")
    else:
        print(from_number)
        return countdown(from_number - 1)
>>> countdown(5)
Sleep for 2 second.
5
Sleep for 2 second.
4
Sleep for 2 second.
3
Sleep for 2 second.
2
Sleep for 2 second.
1
Sleep for 2 second.
Lift off!

Creating Singletons

A singleton is a class with only ONE instance.

The idea of turning a class into a singleton

  • Store the first instance of the class as an attribute
  • Simply freturn the stored instance if later attempts at creating an instance

Example:

import functools

def singleton(cls):
    """Make a class a Singleton class (only one instance)"""
    @functools.wraps(cls)
    def wrapper_singleton(*args, **kwargs):
        if not wrapper_singleton.instance:
            wrapper_singleton.instance = cls(*args, **kwargs)
        return wrapper_singleton.instance

    wrapper_singleton.instance = None
    return wrapper_singleton

@singleton
class TheOne:
    pass
>>> first_one = TheOne()
>>> another_one = TheOne()
>>> id(first_one) == id(another_one)
True
>>> first_one is another_one
True

We can also implement singleton with class decorator:

import functools

class Singleton:
    
    def __init__(self, cls):
        functools.update_wrapper(self, cls)
        self.cls = cls
        self.instance = None

    def __call__(self, *args, **kwargs):
        if not self.instance:
            self.instance = self.cls(*args, **kwargs)

        return self.instance
@Singleton
class OnlyOne:
    pass
>>> first_one = OnlyOne()
>>> another_one = OnlyOne()
>>> id(first_one) == id(another_one)
True
>>> first_one is another_one
True

Caching Return Values

Decorators can provide a nice mechanism for caching and memoization. We can use the function attribute of the inner wrapper function to store a cache lookup table.

Example: Fibonacci

def cache(func):
    """Keep a cache of previous function calls"""
    @functools.wraps(func)
    def wrapper_cache(*args, **kwargs):
        # args_repr = [repr(arg) for arg in args]
        # kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
        # signature = ", ".join(args_repr + kwargs_repr)
        # print(f"{func.__name__}({signature}) cache: {wrapper_cache.cache}")

        cache_key = args + tuple(kwargs.items())
        if cache_key not in wrapper_cache.cache:
            wrapper_cache.cache[cache_key] = func(*args, **kwargs)
        return wrapper_cache.cache[cache_key]
    wrapper_cache.cache = dict()
    return wrapper_cache

@cache
@count_calls
def fibonacci(num):
    if num < 2:
        return num
    return fibonacci(num - 1) + fibonacci(num - 2)

The function attribute wrapper_cache.cache works a s a lookup table, which use the function argument(s) as the key. So now fibonacci() only does the necessary calculations ONCE.

For example, if we have called fibonacci(2) (which returns 1) before, the item {(2, ): 1} will be stored in wrapper_cache.cache, the lookup table. Next time when we call fibonacci(2), we do not need to execute the whole function again. Instead, we just retrieve the value from the lookup table.

>>> fibonacci(10)
Call 1 of 'fibonacci'
...
Call 11 of 'fibonacci'
55

>>> fibonacci(8)
21

Note: in the final call to fibonacci(8), no new calculations were needed, since the eighth Fibonacci number had already been calculated for fibonacci(10).

In the standard library, a Least Recently Used (LRU) cache is available as @functools.lru_cache.

  • has more features
  • You should use @functools.lru_cache instead of writing your own cache decorator
import functools

@functools.lru_cache(maxsize=4)
def fibonacci(num):
    print(f"Calculating fibonacci({num})")
    if num < 2:
        return num
    return fibonacci(num - 1) + fibonacci(num - 2)

Reference