Skip to main content

Learn About Python Decorators by Writing a Function Dispatcher

Python decorators transform how functions work.

@this_is_a_decorator
def some_function():
    return something

"""
>>> some_function()
Something else!
"""

Between the core language and the standard library, there are several decorators that come with Python. Also, many popular frameworks provide decorators to help eliminate boilerplate and make building applications faster and easier.

Decorators are easier to use than to understand.

This article will try to help you understand how decorators work and how to write them. To do that, we'll write a basic implementation of a dispatch function, which will conditionally call function implementations based on the value of an argument.

A Problem in Need of Decorators

How often have you written code that looks something like this?

def if_chain_function(status):
    if status == "red":
        # large block of code
    elif status == "orange":
        # large block of code
    elif status == "yellow":
        # large block of code
    else:
        # some default large block of code

Most of us have written some version of that more than we'd like to admit. Some languages even have a special mechanism (usually called switch) for doing it.

But it has problems.

  • As the number of cases increases, it becomes hard to read.
  • The code blocks in each case will likely evolve independently, some getting refactored into clean functions, some not, and some being a mix of inline code and custom function calls.
  • New cases have to be added explicitly in this function. This precludes pluggable cases from external modules, and also just adds mental load.
  • If the cases are anything other than simple comparisons, the whole thing quickly becomes difficult to reason about.

To solve these problems, making the code more readable and easier to extend, we're going to look at function dispatching.

A little bit about function dispatching

Conventionally, function dispatching is related to the type of the argument. That is, function dispatching is a mechanism for doing different things, depending on whether you pass in an int or a str or a list or whatever.

Python is dynamically typed, so you don't have to specify that a function only accepts some specific type as an argument. But if those different types have to be handled differently, you might end up with code that looks eerily similar to the code above.

def depends_on_type(x):
    if type(x) == str:
        # large block of code
    else if type(x) == int:
        # large block of code
    else if type(x) == list:
        # large block of code
    else:
        # some default large block of code

This has all the same problems mentioned above. But, unlike the first example, Python has a solution to this already:
@functools.singledispatch.

This is a decorator which transforms a function into a single-dispatch generic function. You then register other functions against it, specifying a type of object (that is, a class name). When the function is called, it:

  1. looks up the type of the first argument
  2. checks its registry for that type
  3. executes the function registered for that type
  4. if the type wasn't registered, the original function is executed

If this sounds complicated, don't worry. Using singledispatch is simpler than explaining it.

In [1]:
import functools

@functools.singledispatch
def dispatch_on_type(x):
    # some default logic
    print("I am the default implementation.")
    
@dispatch_on_type.register(str)
def _(x):
    # some stringy logic
    print(f"'{x}' is a string.")
    
@dispatch_on_type.register(int)
def _(x):
    # some integer logic
    print(f"{x} is an integer.")
    
@dispatch_on_type.register(list)
def _(x):
    # some list logic
    print(f"{x} is a list.")
In [2]:
dispatch_on_type(set())
I am the default implementation.
In [3]:
dispatch_on_type("STRING")
'STRING' is a string.
In [4]:
dispatch_on_type(1337)
1337 is an integer.
In [5]:
dispatch_on_type([1,3,3,7])
[1, 3, 3, 7] is a list.

You can do all sorts of cool things with functools.singledispatch, but it doesn't solve the problem in the code at the top of the page. For that, we're going to create a decorator similar to singledispatch that dispatches based on the value of the first argument instead of the type.

Along the way we'll learn more about how decorators work.

Writing Decorators

The @ decorator syntax is syntactic sugar that covers passing a function to another function and then returning a function.

What?

Again, this is easier to show than to explain.

In [6]:
def a_decorator(func):
    return func


# The sweet decorator way...
@a_decorator
def some_function():
    print("Some function.")

# Which has exactly the same effect as...
def some_other_function():
    print("Some other function.")
    
some_other_function = a_decorator(some_other_function)

A decorator is just a function that returns a function.

When used with the @ syntax,

  1. the decorator function is called, with the decorated function passed as an argument;
  2. the return value of the decorator function is assigned to the same name as the decorated function.
  3. When you call the decorated function, you are actually calling the function that was returned by the decorator (which may or may not call the original function's code).

But the above example returned the original function without altering it. The point of decorators is to return something other than the original function, in order to transform the function in some way.

To do this, we usually define another function inside the decorator, and then return that.

In [7]:
def never_two(func):
    def wrapper(*args, **kw):
        x = func(*args, **kw)
        return x if x != 2 else 3
    return wrapper
    
@never_two
def add(x,y):
    return x + y
    
In [8]:
add(1,1)
Out[8]:
3

The wrapper function is defined inside never_two, but it is not executed when never_two is executed (which happens at the line where @never_two appears). Notice — it isn't called anywhere. (That is, you don't see anything like wrapper(1,1).)

Instead, the wrapper function is returned by @never_two, and assigned to the name add. Meanwhile, the code in the original add definition is inside wrapper, where it is called func.

When add(1,1) is called:

  1. The code defined in wrapper is executed, because it was assigned to add.
  2. The arguments passed into add (the two 1s) are passed on to func when it is called at x = func(*args, **kw).
  3. The code originally defined at add (return x + y) is executed, because that code was assigned to func.
  4. The output of func (the original add) is compared to 2, and altered accordingly.
  5. The code defined under wrapper (currently being called add) returns 3.

Two points might be helpful here:

  • Think of a function as just everything from the parenthesis onward, excluding the name. Once you think of a function as a block of code that accepts arguments, which can be assigned to any name, things get a little easier to understand.

  • The (*args, **kw) is a way to collect, pass on, and then unpack all the positional and keyword arguments. A full treatment is beyond the scope of this article. For now, just notice that whatever is passed into wrapper is simply passed on to func.

Writing a Dispatch Decorator

Let's look at the syntax of functools.singledispatch again, and think about what we need to do to emulate it for values instead of types.

In [9]:
@functools.singledispatch
def dispatch_on_type(x):
    print("I am the default implementation.")
    
@dispatch_on_type.register(str)
def _(x):
    print(f"'{x}' is a string.")

Decorators that Return Decorators

Notice that we actually have two decorators:

  • functools.singledispatch
  • dispatch_on_type.register

This means inside singledispatch, the decorated function (in this case, dispatch_on_type) is being assigned an additional attribute, .register, which is also a decorator function.

That might look something like:

In [10]:
def outer_decorator(func):
    
    def inner_decorator(func):
        
        def inner_wrapper():
        
            print("Inner wrapper.")
        
        return inner_wrapper
    
    def wrapper():
        print("The wrapper function.")
    
    wrapper.decorator = inner_decorator
    
    return wrapper

@outer_decorator
def a_function():
    print("Original a_function.") # This will never execute.

@a_function.decorator
def another_function():
    print("Original another_function.") # This will never execute.
In [11]:
a_function()
The wrapper function.
In [12]:
another_function()
Inner wrapper.

Unpacking that a bit:

  • outer_decorator defines two functions,inner_decorator and wrapper
  • wrapper is returned by outer_decorator, so it will be executed when a_function is called
  • inner_decorator is assigned as an attribute of wrapper, so a_function.inner_decorator becomes a usable decorator
  • inner_decorator defines inner_wrapper and returns it, so it will be executed when another_function is called

Decorators with Arguments

You may have noticed that up until now, the decorators created in this article did not included parentheses or arguments when attached to functions. This is because the decorated function is actually passed as the only argument to the function call.

But when registering functions against types, singledispatch included an argument.

In [13]:
@functools.singledispatch
def dispatched():
    return

@dispatched.register(str)
def _():
    return

Incredibly, the way to achieve this is to nest yet another function into the decorator.

That is because, really, register isn't a decorator. Instead, register is a function which returns a decorator when passed an argument.

Let's take our previous example and expand it to include this idea.

In [14]:
def outer_decorator(func):
    
    def faux_decorator_w_arg(arg):
        
        def actual_decorator(func):
            
            def inner_wrapper():
        
                print(f"Inner wrapper. arg was: {arg}")
        
            return inner_wrapper
        
        return actual_decorator
    
    def wrapper():
        print("The wrapper function.")
    
    wrapper.decorator = faux_decorator_w_arg
    
    return wrapper

@outer_decorator
def a_function():
    print("Original a_function.") # This will never execute.

@a_function.decorator("decorator_argument")
def another_function():
    print("Original another_function.") # This will never execute.
In [15]:
a_function()
The wrapper function.
In [16]:
another_function()
Inner wrapper. arg was: decorator_argument

Putting it Together

So now we know how to create decorators that return decorators and accepts arguments. With this, plus a dictionary that maps registered values to functions, we can create a dispatch on value decorator.

In [17]:
def dispatch_on_value(func):
    """
    Value-dispatch function decorator.
    
    Transforms a function into a value-dispatch function,
    which can have different behaviors based on the value of the first argument.
    """
    
    registry = {}

    def dispatch(value):

        try:
            return registry[value]
        except KeyError:
            return func

    def register(value, func=None):
       
        if func is None:
            return lambda f: register(value, f)
        
        registry[value] = func
        
        return func

    def wrapper(*args, **kw):
        return dispatch(args[0])(*args, **kw)

    wrapper.register = register
    wrapper.dispatch = dispatch
    wrapper.registry = registry

    return wrapper
In [18]:
@dispatch_on_value
def react_to_status(status):
    print("Everything's fine.")
    

@react_to_status.register("red")
def _(status):
    # Red status is probably bad.
    # So we need lots of complicated code here to deal with it.
    print("Status is red.")
In [19]:
react_to_status("red")
Status is red.

There are a few things here which might not be obvious. So let's take a closer look.

def dispatch(value):

        try:
            return registry[value]
        except KeyError:
            return func

This is called by wrapper, and is the mechanism that determines which registered function is executed. It looks in the registry and returns the appropriate function (without executing it). If the value isn't registered, this will raise a KeyError. The except block catches that error and returns the original function.

def register(value, func=None):

        if func is None:
            return lambda f: register(value, f)

        registry[value] = func

        return func

This acts as both the faux_decorator and the actual_decorator.

It can be called with one or two positional arguments; if the second one is omitted, it is set to None.

At @react_to_status.register("red"), it is being called with only the value argument. This causes the lambda expression to be returned, with value already interpreted. (That is, the return value is lambda f: register("red", f)).

This is the same as:

if func is None:
        def actual_decorator(func):
            register(value, func)
        return actual decorator

But the lambda expression is a bit easier to read, once you know what it is doing.

This returned lambda function is then the actual decorator, and is executed with the wrapped function as its one argument. The function argument is them passed to register, along with the value that got set when the lambda was created.

Now register runs again, but this time it has both arguments. The if func is None is skipped, and the function is added to the registry with value as the key. The function is returned back to the point when the register decorator was called, but it gets assigned to the name _, because we never need to call it directly.

def wrapper(*args, **kw):
        return dispatch(args[0])(*args, **kw)

This is the function that actually gets executed when react_to_status is called. It calls dispatch with the first argument (arg[0]), which return the appropriate function. The returned function is immediatly called, with *args, **kw passed in. Any output from the function is returned to the caller of react_to_status, which completes the entire dispatch process.

Going Further

This tutorial looked at value dispatch in order to dig into how decorators work. It does not provide a complete implementation for a practical value dispatch decorator.

For example, in practice you'd probably want value dispatch to include:

  • values within a range
  • values in a collection
  • values matching a regular expression
  • values meeting criteria defined in a fitness or sorting function

And we didn't even talk about additional functools features that help sort out introspection or dealing with the many problems created by decorators.

For a more complete, production ready implementation of this idea, see Dispatch on Value by minimind.


Credits

The final form of dispatch_on_value was based heavily on the Ouroboros implementation of functools.singledispatch.

Comments

Comments powered by Disqus