Stupid Python Tricks: Learn the ABCs (and SDFs)

My favorite language is Python. I just love digging into all of the weird and esoteric parts of it. But since my job is primarily Ruby, and I love my job enough to throw my hobby programming at it, I don’t spend as much time getting better at Python as I’d like to. That, plus the fact that a lot of people haven’t played around with the new Py3 features, means I think it’d be fun to share a few stupid tricks I’ve discovered. So let’s talk about functions!

Setup

Define a “point” as a named tuple with fields x and y. We want to find the distance from the origin for a given point. Here’s one way to do this:

from collections import namedtuple
from math import sqrt
Point = namedtuple('Point', 'x y’)
def distance(p):
    return sqrt(p.x**2 + p.y**2)

Now let’s add polar coordinates. Instead of x and y, we have r and theta. How would we update our function? The normal way would be to add a conditional in the body of distance. So let’s do that:

from collections import namedtuple
from math import sqrt
CartesianPoint = namedtuple('CartesianPoint', ['x', 'y'])
PolarPoint = namedtuple('PolarPoint', ['r', 'theta'])
def distance(p):
    if isinstance(p, CartesianPoint):
        return p.x + p.y
    elif isinstance(p, PolarPoint):
        return p.r
    else:
        raise TypeError

Heck, now let’s add a third dimension and drop in 3D Cartesian coordinates, cylindrical coordinates, and spherical coordinates.

from collections import namedtuple
from math import sqrt
CartesianPoint = namedtuple('CartesianPoint', 'x y')
PolarPoint = namedtuple('PolarPoint', 'r theta')
CartesianPoint3D = namedtuple('CartesianPoint3D', 'x y z')
CylindricalPoint = namedtuple('CylindricalPoint', 'rho phi z')
SphericalPoint = namedtuple('SphericalPoint', 'r phi theta')
def distance(p):
    if isinstance(p, CartesianPoint):
        return sqrt(p.x**2 + p.y**2)
    elif isinstance(p, PolarPoint):
        return p.r
    elif isinstance(p, CartesianPoint3D):
        return sqrt(p.x**2 + p.y**2 + p.z**2)
    elif isinstance(p, CylindricalPoint):
        return sqrt(p.rho**2 + p.z**2)
    elif isinstance(p, SphericalPoint):
        return p.r
    else:
        raise TypeError

This will not scale. Normally people would say “so use a class already”, which is a valid path that has its own pros and cons. For now, though, it’d be nice to use a purely functional solution, and we’ll look for that instead.

Single-Dispatch Functions to the rescue

Single-dispatch functions were added in PEP 443 to help with exactly this kind of problem. We define a single function for the general case, and then “register” a new instance of that function for each special type it should handle. One example is a function that works for both a number and a sequence of numbers. Here, each of our points is a distinct type, so instead of defining one function with six branches, we define six dispatch functions:

from functools import singledispatch
from collections import namedtuple
from math import sqrt
CartesianPoint = namedtuple('CartesianPoint', ['x', 'y'])
PolarPoint = namedtuple('PolarPoint', ['r', 'theta'])
CartesianPoint3D = namedtuple('CartesianPoint3D', ['x', 'y', 'z'])
CylindricalPoint = namedtuple('CylindricalPoint', 'rho phi z')
SphericalPoint = namedtuple('SphericalPoint', 'r phi theta')
@singledispatch
def distance(p):
    raise TypeError
@distance.register(CartesianPoint)
def _(p):
    return sqrt(p.x**2 + p.y**2)
@distance.register(PolarPoint)
def _(p):
    return p.r
@distance.register(CartesianPoint3D)
def _(p):
    return sqrt(p.x**2 + p.y**2 + p.z**2)
@distance.register(CylindricalPoint)
def _(p):
    return sqrt(p.rho**2 + p.z**2)
@distance.register(SphericalPoint)
def _(p):
    return p.r

This might seem even more complicated than the if-else chain, but it has a few advantages. One, it’s easier to extend. If we want to add another coordinate type, we don’t have to modify the existing function; we just register a new one. It makes it much more explicit what we need to test. Instead of trying to see what lines our test cover, we know there are six specific functions we need to check. And since we have multiple independent functions, we can slap different decorators on each one. If the cartesian calculations were a bottleneck, we could cache those without also having to cache the polar ones. Pretty neat!

You might notice that two of the six functions are basically duplicates. It would be nice if we could simplify this further and only need four functions. And we can. But to do that, we have to get weird.

Abstract Base Classes

Abstract Base Classes (ABCs) were added in PEP 3119 as a supplement to duck typing. One of the reasons duck typing works so well is that you don’t have to care about what an object is, just what it can do. It shouldn’t matter whether you pass strings or integers into max, as long as its orderable. This works great, but it doesn’t scale too well. What if your function requires your object to be orderable and iterable and callable?

ABCs provide two tools to help with that. First, you can define abstract methods. Classes that inherit the ABC (concrete subclasses) can’t be instantiated unless you’ve defined all of those methods. Pretty similar to how ABCs work in C++ or interfaces work in Java.

The second tool is much stranger. By defining __subclasshook__ on your ABC, any class that passes the check is considered a “virtual” subclass of the ABC. That allows you to easily extend it to third party objects. It also means that we can define the conditions for a point to be “Cartesian” or “Polar” and use that to “simplify” our code:

from abc import ABC
from functools import singledispatch
from collections import namedtuple
from math import sqrt
class NCartesianPoint(ABC):
    @classmethod
    def __subclasshook__(cls, point):
        return "CartesianPoint" in point.__name__
class NSpherePoint(ABC):
    @classmethod
    def __subclasshook__(cls, point):
        if 'Point' in point.__name__:
            return 'r' in point._fields
        return NotImplemented

CartesianPoint = namedtuple('CartesianPoint', 'x y')
PolarPoint = namedtuple('PolarPoint', 'r theta')
CartesianPoint3D = namedtuple('CartesianPoint3D', 'x y z')
CylindricalPoint = namedtuple('CylindricalPoint', 'rho phi z')
SphericalPoint = namedtuple('SphericalPoint', 'r phi theta')
@singledispatch
def distance(p):
    raise TypeError
@distance.register(NCartesianPoint)
def _(p):
    return sqrt(sum((x**2 for x in p)))
@distance.register(NSpherePoint)
def _(p):
    return p.r
@distance.register(CylindricalPoint)
def _(p):
    return sqrt(p.rho**2 + p.z**2)

Now if we want to add another coordinate type, we might not even need to register the function. It’ll work for a 4-Sphere, a 37-Sphere, etc. Sure, the code is more complex and arcane, but there’s a reason this isn’t called “smart and reasonable python tricks”.

Caveat: Unfortunately, you can’t use __subclasshook__ with gradual typing and mypy, part because the parser is fairly simple and part because adding it means solving the halting problem.

Summary

• Single-dispatch functions let you overload a function based on argument type.
• Abstract Base Classes let you (among other things) dynamically assign additional types to arbitrary classes.
• Combining the two makes your code weird and magical.
• It also makes your type-checker sad, so use sparingly.

Further Reading

• Single-dispatch functions and PEP 443
• Abstract Base Classes and PEP 3119
• Some more theory on SDFs (and how to use with enums!)
• Why use Abstract Base Classes in Python?