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.
Leave a comment