Learn Type Variables and Generic Types in Python with Simple Examples

Image by kirillslov on Pixabay

In this post, we will introduce the basics of type variables and generic functions and classes with simple code examples. Using type variables and generic types can make our code more robust and easier to read.

Even if we may not have many chances to use them directly in our own code, it’s still very helpful to understand the basics of generic types so we can read the source code or documentation of some libraries more easily. A prominent example is Pydantic which uses generic types very commonly.


Install Python and mypy

In order to use the latest features of Python typing, the newest version of Python will be installed in a conda virtual environment. You should have at least Python version 3.10 in order to follow the examples. Mypy is also installed to check the typing and reveal the errors.

conda create -c conda-forge -n typing python=3.11 mypy=1.5

conda activate typing

Note that we need to set the channel to conda-forge for conda in order to install the latest version of mypy with conda.


Type variables

A type variable is a special variable in Python used to represent a type. We can define a type variable using the TypeVar construct from the typing module:

from typing import TypeVar

T = TypeVar("T")

Note that the name of the type (the first argument of TypeVar) must be the same as the variable name, otherwise mypy would raise an error.

By default, a type variable is unconstrained and can be any type, for example:

def echo(a: T):
    print(a)

echo(123)   # OK
echo(True)  # OK
echo("abc") # OK

Now let’s see a slightly more complex example:

def add(a: T, b: T) -> T:
    return a + b

If you use mypy to check this file, you will see this error:

error: Unsupported left operand type for + ("T")  [operator]
Found 1 error in 1 file (checked 1 source file)

This is because by default the T type variable can be any type, but not all types support the + operator.

To solve this problem, we can constrain the allowed types of T to some specific ones:

T = TypeVar("T", int, float)

def add(a: T, b: T) -> T:
    return a + b

This time mypy will not complain about the types.

Besides specifying the allowed types explicitly, we can set an upper bound for the type variable, which means that the actual type must be a subtype of the bound type:

from typing import TypeVar

class Animal:
    pass

class Horse(Animal):
    pass

class Robot:
    pass


T = TypeVar("T", bound=Animal)

def lets_move(animal: T):
    pass

lets_move(Animal())  # OK
lets_move(Horse())   # OK
lets_move(Robot())   # Incompatible type

Generic function

A generic function is a function that can be parameterized by one or more types, which means the function can operate on data of various types.

The most common use case is to pass a list/tuple of elements of various types to a function. The type variables in the generic functions can establish a relationship between argument types and return types.

from typing import TypeVar, Sequence

T = TypeVar("T")

def first_of_list(l: list[T]) -> T:
    return l[0]

def first_of_tuple(t: tuple[T, ...]) -> T:
    return t[0]

def first_of_sequence(s: Sequence[T]) -> T:
    return s[0]

l = [1 ,2, 3]
t = ("a", "b", "c")

l0 = first_of_list(l)   # l0 is of type int.
t0 = first_of_tuple(t)  # s0 is of type str.

first_of_list(t)    # Error: type incompatible
first_of_tuple(l)   # Error: type incompatible

first_of_sequence(l)
first_of_sequence(t)

Besides, we can see that Sequence is a handy container type that is compatible with both lists and tuples, it can be used for any type that supports the __getitem__ and __init__ magical method such as list, tuple and str.

We can pass more than one type variable to a generic function, which is most common when some dictionary type is used in the argument or return value.

Let’s see a simple example function that takes a dict as the argument and return a tuple of two sets that contain the keys and unique values, respectively:

from typing import TypeVar

K = TypeVar("K")
V = TypeVar("V")


def get_keys_and_values(d: dict[K, V]) -> tuple[set[K], set[V]]:
    return set(d.keys()), set(d.values())

x = get_keys_and_values({1: 100, 2: 200})
# type of x is: tuple[set[int], set[int]]

y = get_keys_and_values({"a": True, "b": False})
# type of y is:tuple[set[str], set[bool]]

Generic classes

Similarly, a generic class is a class that can be parameterized with one or more types. A generic class needs to inherit from the Generic construct:

from __future__ import annotations

from typing import Generic, TypeVar

T = TypeVar("T")


class Node(Generic[T]):
    def __init__(self, data: T) -> None:
        self.data = data
        self.children: list[Node] = []

    def add_child(self, child_node: Node):
        self.children.append(child_node)

Note that we need to import annotations from __future__ to support using the class itself as a type in its own definition, which is called forward reference.

To create instances of a class with uniform types, we need to use type aliases, which are defined by simple variable assignments:

IntNode = Node[int]
node1 = IntNode(100)    # OK
node1 = IntNode("abc")  # Incompatabile type

StrNode = Node[str]
node3 = StrNode(100)    # Incompatabile type
node4 = StrNode("abc") # OK       

We can also pass more than one type variable to a generic class:

from __future__ import annotations

from typing import Generic, TypeVar

L = TypeVar("L")
V = TypeVar("V")

class Node(Generic[L, V]):
    def __init__(self, label: L, data: V) -> None:
        self.label = label
        self.data = data
        self.children: list[Node] = []

    def add_child(self, child_node: Node) -> None:
        self.children.append(child_node)

IntNode = Node[int, int]
node1 = IntNode(1, 100)   # OK
node1 = IntNode(1, "abc") # Incompatabile type

StrNode = Node[str, str]
node3 = StrNode("a", "value")  # OK
node4 = StrNode(1, 200)        # Incompatabile type

In this post, we introduced the basics of type variables and generic functions and classes. We may not have many chances to use them directly in our own code, however, it’s very helpful to understand the basics in order to read the source code of some libraries. A prominent example is Pyantic 2.0 which uses generic types very commonly.


Related articles:



Leave a comment

Blog at WordPress.com.