@singledispatch in Python - Get Rid of Those Long If-Statements

In this post I’d like to cover a cool python decorator that’s part of the built-in functools library called @singledispatch. It’s a favorite of mine and something I’ve been using quite a bit over the past year. I’m surprised how little I see it in projects I come across so hopefully this post will help it find its way into a repo or two. ;)

To understand @singledispatch, you have to understand the pre-existing alternative that it’s meant to replace. Well, to really understand it you should read through PEP 443, the proposal that led to its addition to the standard library. But let’s say you’d rather read this post instead! At the root is a simple problem. You have a function that does something generic and how it actually does that “something” depends on the type of input it’s given. Let’s think of a super simple case of a function that names things. When you run the function, it hits a REST endpoint that returns a random name. I’m sure there’s a public API out there we can use, but for this post, we’ll just use simple lists and the built-in random.

import random

names = ["John", "Jacob", "Jingleheimer", "Schmidt"]

def nameit():
    return random.choice(names)

Great! A simple function for a simple usecase. Now say the API is pretty clever and allows you to narrow the selection pool to a certain type of name. To keep it simple, we’ll imagine the only options besides a human name is a superhero, pet, or robot name.

names = ["John", "Jacob", "Jingleheimer", "Schmidt"]
superhero_names = ["Captain Crunch", "Fighting Frito", "Big Tuna", "The Math Whisperer"]
pet_names = ["Fluffy", "Jet", "Rocko", "Popper"]
robot_names = range(10000, 99999)

These perfectly match to constructs in your app; In other words, thing is some instance of a SuperHero, Pet, or Robot class defined in your application. There’s also a Gadget class but that’s ok, you’ll just group that under the robot naming logic which calls the robot endpoint and returns a random number.

A simple way to narrow in the name selection would be to accept the object as an argument, check its type, and then call the appropriate endpoint.

class SuperHero:
    ...

class Pet:
    ...

class Robot:
    ...

class Gadget:
    ...

def nameit(thing):
    """Get a random name for a thing

    `thing` can be an instance of SuperHero, Pet, Robot, or Gadget. This function
    returns a string for all cases except for Robot or Gadget instances, in which
    case it returns an integer.
    """
    if isinstance(thing, SuperHero):
        return random.choice(superhero_names)
    elif isinstance(thing, Pet):
        return random.choice(pet_names)
    elif isinstance(thing, (Robot, Gadget)):
        return random.choice(robot_names)
    
    # If not a known type, choose a name from the `names` list
    return random.choice(names)

Alright, this works nicely but the code is starting to look a bit untidy. The series of if-statements and multiple requests to different routes causes you to have to briefly squint to follow what’s happening. Furthermore, the docs require you to mentally process the context in which you’ll be using the function. It returns a string if you pass in a SuperHero or Pet but it returns an integer if you pass in a Robot or Gadget. Even though you can separate these out into different functions, like name_a_superhero(x) or name_a_pet(x), you don’t want to complicate the usability that comes with a generic function that allows users to just call nameit(x). This is where @singledispatch can be useful.

With the @singledispatch decorator, you can define multiple variations of a single function and have the right variation used depending on the type of the first argument. The function that’s dispatched is based on a single argument (the first argument), hence the name of the decorator. To convert the nameit function into a single dispatch function, simply add the decorator.

from functools import singledispatch

@singledispatch
def nameit(thing):
    """Get a random name for a thing

    `thing` can be an instance of SuperHero, Pet, Robot, or Gadget. This function
    returns a string for all cases except for Robot or Gadget instances, in which
    case it returns an integer.
    """
    if isinstance(thing, SuperHero):
        return random.choice(superhero_names)
    elif isinstance(thing, Pet):
        return random.choice(pet_names)
    elif isinstance(thing, (Robot, Gadget)):
        return random.choice(robot_names)
    
    # If not a known type, choose a name from the `names` list
    return random.choice(names)

Now nameit is a single dispatch function. This means the function itself can be used as a decorator to register other type-based variations of itself. Let’s use @nameit to register a function to handle cases where thing is a SuperHero instance.

@nameit.register(SuperHero)
def _name_a_superhero(superhero) -> str:
    """Get a random super hero name"""
    return random.choice(superhero_names)

This new function _name_a_superhero is now registered to the nameit function and gets called automatically whenever the first argument to nameit is an instance of SuperHero. Since the function is never called directly, its name doesn’t actually matter and a common convention is to just use an underscore.

@nameit.register(SuperHero)
def _(superhero) -> str:
    """Get a random super hero name"""
    return random.choice(superhero_names)

Let’s create and register a function for Pet and Robot instances.

@nameit.register(Pet)
def _(pet) -> str:
    """Get a random pet name"""
    return random.choice(pet_names)

@nameit.register(Robot)
def _(robot) -> int:
    """Get a random robot name"""
    return random.choice(robot_names)

Another cool thing about single dispatch functions is that if you have multiple types that should be treated the same, you can simply stack the register decorators on the corresponding function. Since we name Gadget instances the same way we do Robot instances, we can stack that right on top of the same function!

@nameit.register(Gadget)
@nameit.register(Robot)
def _(robot) -> int:
    """Get a random robot name"""
    return random.choice(robot_names)

So what about cases where the type of the first argument has no type-based function registered? In that case, single dispatch falls back to the original nameit function. It’s common to raise a TypeError or NotImplementedError here, but in our case we just want to return a random name from the names list.

@singledispatch
def nameit(thing) -> str:
    """Get a random name for a thing"""
    return random.choice(names)

Here’s our code so far.

from functools import singledispatch
import random

class SuperHero:
    ...

class Pet:
    ...

class Robot:
    ...

class Gadget:
    ...

names = ["John", "Jacob", "Jingleheimer", "Schmidt"]
superhero_names = ["Captain Crunch", "Fighting Frito", "Big Tuna", "The Math Whisperer"]
pet_names = ["Fluffy", "Jet", "Rocko", "Popper"]
robot_names = range(10000, 99999)

@singledispatch
def nameit(thing) -> str:
    """Get a random name for a thing"""
    return random.choice(names)

@nameit.register(SuperHero)
def _(superhero) -> str:
    """Get a random super hero name"""
    return random.choice(superhero_names)

@nameit.register(Pet)
def _(pet) -> str:
    """Get a random pet name"""
    return random.choice(pet_names)

@nameit.register(Gadget)
@nameit.register(Robot)
def _(robot) -> int:
    """Get a random robot name"""
    return random.choice(robot_names)

The flurry of if-statements have been replaced with a set of small and concise functions with focused docstrings, all while maintaining the single entrypoint of calling the nameit function and passing in an object that you want to name.

nameit(SuperHero())  # Big Tuna
nameit(Pet())  # Rocko
nameit(Robot())  # 26632
nameit(Gadget())  # 78444
nameit(0)  # Schmidt

This example converted what was already a pretty simple function into a single dispatch function to illustrate how the decorator works. For more complicated functions, however, it’s even more apparent how single dispatch lets your code breathe by breaking up convoluted conditionals and calls to other functions and placing each flow of the conditional logic into its own function. That means each of these code flows can have their own docstring, typehints, and general area in your codebase. They can even live in different source code files! Another thing is that in some cases you can experience a significant performance improvement when you replace slower if-else conditionals with a single dispatch that’s more akin to a much faster dictionary lookup.

I hope you find this helpful and if you have any thoughts or questions, or if you’ve tried this in one of your projects after reading this post and want to let me know how it went, feel free to reach out on twitter @samuelredai, message me on LinkedIn, or send me an email at samuelspersonalemail@gmail.com.

BONUS:

If one of your registered functions won’t be dispatched for more than one type, you can simplify the decorator by letting the type be inferred from the typehint of the first argument to the function.

So instead of…

@nameit.register(SuperHero)
def _(superhero) -> str:
    """Get a random super hero name"""
    return random.choice(superhero_names)

You can use…

@nameit.register
def _(superhero: SuperHero) -> str:
    """Get a random super hero name"""
    return random.choice(superhero_names)

note: Inferring the type from the typehint was added in Python 3.7