Type Hinting: Union and Sum types

Today we will take a look at an interesting feature in Python type hints: Union types (also sometimes called as sum types in functional programming languages).

In this article, we will be working with an example of an e-commerce store that sells books and video games. To represent that data, we create the following two data classes

from dataclasses import dataclass

@dataclass(frozen=True)
class Book:
    title: str
    author: str
    isbn: str
    price: float
    
@dataclass(frozen=True)
class VideoGame:
    title: str
    developer: str
    platform: str
    price: float

Note that these two classes are independent. They don't inherit from each other, or from any base class.

Creating union types

An item in this shop may be a book, or it might be a video game. We can represent that by using a union. A union takes a list of types and the variable can be assigned to objects of any of those types

from typing import Union

# ✅ OK
item1: Union[Book, VideoGame] = Book('Foundation', 'Asimov', '0380009145', 800)

# ✅ OK
item2: Book | VideoGame = VideoGame('Super Mario Odyssey', 'Nintendo', 'Switch', 4000)

Here, we use the type Union[Book, VideoGame] to represent a type that could be a Book or a VideoGame. Such types are called Union types or Sum types.

In the second line above, we use the syntax Book | VideoGame. This is another, shorter way of representing a union type that was introduced in Python 3.10.

Now, it can be quite irritating writing Union[Book, VideoGame] over and over again whenever we use a variable to represent an item. To solve this problem, Python allows us to create type aliases.

Type aliases are type variables that represent a more complicated type. For example, we can do this

from typing import TypeAlias

Item: TypeAlias = Book | VideoGame

Now that we have this alias, we can use the Item type everywhere. mypy will know that it means Book | VideoGame.

# ✅ OK
item3: Item = Book('2001: A Space Odyssey', 'Clarke', '9780451457998', 800)

Working with union types

With sum types, the variable could be any of those types and our code needs to handle that. The function below won't work because it item is a Book, it won't have a platform field, and if it's a videogame it won't have author.

def print_item(item: Item):
    print(item.title) # ✅ OK
    print(item.price) # ✅ OK
    print(item.author) # ❌ VideoGame doesn't have author
    print(item.platform) ❌ Book doesn't have platform

mypy will catch this and flag an error on both lines

The right way is to do an isinstance check and handle the two variants separately. mypy will not give any error for the code below. mypy can figure out that inside the isinstance condition the item must be a Book and therefore in the else clause it has to be a VideoGame

def print_item(item: Item):
    print(item.title) # ✅ OK
    print(item.price) # ✅ OK
    if isinstance(item, Book):
        # ✅ mypy knows if we get here then item is a Book
        print(item.author)
    else:
        # ✅ mypy knows its not a Book so it must be a VideoGame
        print(item.platform)

Once we have everything set up, we can easily create a cart of items. The items could be books or video games, but nothing else.

# list containing [Book, VideoGame, Book]
cart: list[Item] = [item1, item2, item3]

for item in cart:
    print_item(item)

We can loop over it and call the print_item functions and everything works properly.

Now imagine that later on we start selling music. So we create a class to represent that, and update the Item alias to include Music

@dataclass(frozen=True)
class Music:
    title: str
    performer: str
    label: str
    price: float

Item: TypeAlias = Book | VideoGame | Music

Immediately, mypy will flag an error in the print_item function.

def print_item(item: Item):
    print(item.title) # ✅ OK
    print(item.price) # ✅ OK
    if isinstance(item, Book):
        # ✅ mypy knows if we get here then item is a Book
        print(item.author)
    else:
        # ❌ now item could be VideoGame or Music in the else part
        print(item.platform)

Once we fix the issue, we can add music into our cart.

def print_item(item: Item):
    print(item.title) # ✅ OK
    print(item.price) # ✅ OK
    if isinstance(item, Book):
        # ✅ mypy knows if we get here then item is a Book
        print(item.author)
    else if isinstance(item, VideoGame):
        # ✅ mypy knows if we get here then item is a VideoGame
        print(item.platform)
    else:
        # ✅ if we get here then item nust be Music
        print(item.performer)

In a large codebase, it is easy to forget to update the code in some place, and this could lead to bugs. mypy is clever to navigate the isinstance checks and deduce what the possible data types should be. It then helpfully points out when there is a bug in the code.

Summary

So that is all about union types, also known as sum types.

Using union types makes it possible for our code to accept arguments in many forms. Maybe you have a function that takes an input which can be a filename or a file object. You can easily represent this type of parameter using a union type and mypy will happily validate that the type is not violated by any of the callers.

It is a very powerful feature, especially in a dynamic language like Python where these kinds of functions are very common.