keyboard-shortcut
d

typevars in python 🐍

12min read

an image

I think I need a TypeVar?

My first piece of advice is, "if you need your function to be generic, ask just how generic it needs to be?". If you know what types you are expecting, you might just save yourself a lot of hassle by avoiding TypeVars and being quite specific in your requirements (and tackling flexibility when it is needed later).

If you're confident that your typehints might benefit from a TypeVar, read on...

What is a TypeVar?

TypeVar is used in type hinting to declare a variable that can represent multiple types. TypeVar allows functions and classes to be flexible in the type of objects they process but then be consistent with that type later on.

The way I think of it is "whatever type is inferred on the way in, is the type used on the later on". If you provide a string on the way in, it will be a string on the way out. If you provide an integer on the way in, it will be an integer on the way out. If you provide a list with a typevar and that list contains a mixture of integers, floats, strings, etc... then on the way out it will be a list of "objects" (as that was inferred on the way in).

An example

Below is a function that takes a random sample from a list. It makes a random decision whether or not to keep an elemen in the list. It doesn't really matter what is in the list. This is function accepts a list of elements of "generic" type.

One common way of accepting "any" type is using the Any type hint, like so:

import random
from typing import Any


def sample(population: list[Any]) -> list[Any]:
    """
    Return a sample of elements from the full population.
    """
    return [element for element in population if random.choice([True, False])]

The problem with Any

The loss of information

Using this function, a few problems arise because of just how vague Any is.

If I use my sample function on a list of integers integer_list = [1,2,3,4,5] my type checker is capable of identifying that this is of type list[int]. This is good 👍.

The problem is that when I spit out my new subset of the list after running it through sample the type checker is now completely incapable of identifying what this type is because our function signature says it is a list of Any. Static type checkers contain no context for the elements in that list because of the vague function signature. This is bad 👎 (for your linter and your IDE).

We can demonstrate this using mypy and its reveal_type magic. We can quickly run mypy with no setup using uv by running uvx run mypy <name_of_file>.py

import random
from typing import Any


def sample(population: list[Any]) -> list[Any]:
    """
    Return a sample of elements from the full population.
    """
    return [element for element in population if random.choice([True, False])]

integer_list = [1, 2, 3, 4, 5]
reveal_type(integer_list)
integer_sample = sample(integer_list)
reveal_type(integer_sample)

We can clearly see where information is lost:

13: note: Revealed type is "builtins.list[builtins.int]"
15: note: Revealed type is "builtins.list[Any]"

Any being a fine input to anything

Because of the vague nature of Any, this loss of information could cause problems if we connect sample to another function. Our type checker would be incapable of identifying the bug meaning we would only capture this issue at runtime.

import random
from typing import Any


def sample(population: list[Any]) -> list[Any]:
    """
    Return a sample of elements from the full population.
    """
    return [element for element in population if random.choice([True, False])]

string_list = ["a", "b", "c"]  # a list of strings
string_sample = sample(string_list)  # a (possibly) smaller list of strings (or maybe an empty string, but that's a separate issue...)
sum(string_sample)  # `sum` will fail at runtime because it can't add strings :(

However, running our type checker, we don't get any issues.

Success: no issues found in 1 source file

This is because Any is basically a wildcard in type hint land. Any is "unconstrained". A static type checker will treat every type as being compatible with Any and Any as being compatible with every type

Avoiding the problem without using TypeVars

If we have a specific use-case in mind, I would recommend being more specific in our sample function:

def sample_numbers(number_population: list[Union[int, float]]) -> list[Union[int, float]]:
    """
    Return a sample of numbers from the full population of numbers.
    """
    return [numbers for numbers in population if random.choice([True, False])]

This would correctly throw lots of errors in the previous example where we provided strings to these functions before throwing it into a function that only accepts numbers. This works perfectly for our use case. Unfortunately, it is only suitable for our one specific use-case.

Solving the problem with TypeVars

If we want sample to be "generic" we need to use TypeVars. Improving our function with TypeVar will enable us to:

  • ✅ Understand the code better. It is more obvious that the type for our argument is the same type as the return statement.
  • ✅ Preserve information. Whatever the type checker or linter sees going into the function is what gets returned.

So, how do we do it?

import random
from typing import TypeVar

T = TypeVar("T")

def sample(population: list[T]) -> list[T]:
    """
    Return a sample from the full population.
    """
    return [numbers for numbers in population if random.choice([True, False])]

integer_list = [1, 2, 3, 4, 5]
reveal_type(integer_list)
integer_sample = sample(integer_list)
reveal_type(integer_sample)

string_list = ["a", "b", "c", "d", "e"]
reveal_type(string_list)
string_sample = sample(string_list)
reveal_type(string_sample)

This allows us to preserve information in our type checker. As a result, we get an appropriate error when we try to sum the string_sample:

import random
from typing import TypeVar

T = TypeVar("T")

def sample(population: list[T]) -> list[T]:
    """
    Return a sample from the full population.
    """
    return [element for element in population if random.choice([True, False])]

string_list = ["a", "b", "c", "d", "e"]
reveal_type(string_list)
string_sample = sample(string_list)
reveal_type(string_sample)

sum(string_sample)

mypy gives us an error (albeit, not a particularly helpful one...):

test.py:13: note: Revealed type is "builtins.list[builtins.str]"
test.py:15: note: Revealed type is "builtins.list[builtins.str]"
test.py:17: error: Argument 1 to "sum" has incompatible type "list[str]"; expected "Iterable[bool]"  [arg-type]

This solves most of our needs!

Mixing types...

All of the examples so far have had clear uniform types. The elements in the list have all been the same type. What happens if we mix these? How do these get inferred?

mixed_example = [1, "a", 3.0]
reveal_type(mixed_example)

It would be amazing if this was correctly inferred as list[Union[int, str, float[]], but the news isn't great...

test.py:2: note: Revealed type is "builtins.list[builtins.object]"

The object type only allows "universally compatible" operations to occur. Operations like print, evaluating truthiness and type checking are acceptable but "specific" operations like adding, multiplying, strip (for a string) are unacceptable. I've pulled together some examples from the mypy docs:

def try_things_on_my_object(obj: object) -> None:
    truthy = obj == True  # ✅
    print(obj)  # ✅
    print(isinstance(obj, int))  # ✅
    obj * 2  # ❌
    obj + 1  # ❌
    obj - 1  # ❌
    open(obj)  # ❌

Fix the mix - constraints

We can ensure no "mixing" while still allow "generality" by adding constraints. Adding constraints to a TypeVar means that all instances must match exactly one of the types provided (the specific type all elements match doesn't matter, hence keeping this function "generic").

For example, the below TypeVar has the constraints str and int which effectively checks that exactly one of the conditions below is met: - all elements are strings - all elements are floats

import random
from typing import TypeVar

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

def sample(population: list[T]) -> list[T]:
    """
    Return a sample from the full population.
    """
    return [element for element in population if random.choice([True, False])]

This keeps our function generic (you can still apply it to multiple types so long as they are consistent).

It's worth noting that you have to provide a minimum of two constraints to get this working. This makes sense given that the whole point in the first place was to allow multiple types.

There are a few things to be aware of though...

  • int and float can be mixed at will... This is a bit annoying because list[int] and list[float] aren't the same thing (and your type checker will generally complain about this), but when using constraints we can mix them freely.
  • If you mention a type, you can freely mix subtypes. For example, if you define class Parent, then two child classes Son(Parent) and Daughter(Parent), then include Parent in your constraints you can still freely mix Son and Daughter wherever your TypeVar is used.
  • bool is a subtype of int, so, be careful! Again, you can mix bool and int.
  • You can provide the same constraint twice, meaning that basically there is only one constraint 🤷

Fix the mix - bound

There is also bound. This exists for when you want to limit your "generic" type to one specific type, but you still want tools to be able to infer the exact subtype. Again, it all comes back to preserving information as types pass throuh a system.

This one is a bit easier to demonstrate with an example.

Imagine you define a series of classes that use inheritance.

from copy import copy


class Animal: ...


class HousePet(Animal):
    def speak(self) -> str:
        raise NotImplementedError


class Dog(HousePet):
    def speak(self) -> str:
        return "woof!"


class Cat(HousePet):
    def speak(self) -> str:
        return "meow!"


def clone_my_pet(pet: HousePet) -> HousePet:
    new_pet = copy(pet)
    return new_pet


def send_to_doggy_day_care(dog: Dog) -> None:
    # transfer_lots_of_money()
    # something else...
    return None


fido = Dog()

frankenfido = clone_my_pet(fido)

send_to_doggy_day_care(fido)
send_to_doggy_day_care(frankenfido)

Here I've defined some basic classes that use inheritance, but then I have defined some methods that only work on specific subtypes. clone_my_pet only works on HousePets and send_to_doggy_day_care only works on Dogs. This is fine, until you start chaining some of these types together... Running this code through our type checker informs us that we can't clone our dog and send them to day care. This is because information is lost when cloning. Instead, we need use a TypeVar to preserve information.

from copy import copy
from typing import TypeVar


class Animal: ...


class HousePet(Animal):
    def speak(self) -> str:
        raise NotImplementedError


class Dog(HousePet):
    def speak(self) -> str:
        return "woof!"


class Cat(HousePet):
    def speak(self) -> str:
        return "meow!"


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


def clone_my_pet(pet: T) -> T:
    new_pet = copy(pet)
    return new_pet


def send_to_doggy_day_care(dog: Dog) -> None:
    # transfer_lots_of_money()
    # something else...
    return None


fido = Dog()

frankenfido = clone_my_pet(fido)

send_to_doggy_day_care(fido)
send_to_doggy_day_care(frankenfido)

Using a TypeVar here allows us to preserve information. Using bound=HousePet restricts this function from accepting anything that isn't a HousePet. This gives us the functionality we need. The difference between this and using a HousePet constraint is that our type checker is able to infer the exact subtype (in this case Dog) that is provided when we call send_to_doggy_day_care(frankenfido). If we set one of the constraints as HousePet, we would get the same error as before.

The constraints mechanism only allows us to infer the exact types provided. The bound method represents an "upper bound" of the acceptable type, but lets us infer the exact subtype.

Things to take away

  • Avoid TypeVars if being generic is "nice" but not strictly necessary. Save yourself the testing by just being more specific.
  • If you use TypeVar do lots of testing and check what your type checker is inferring (using functionality simmilar to mypy's reveal_type).
  • If you see that mypy has inferred something as object, this is probably undesirable and mypy has just seen more types than it should...
  • Be careful of "special" cases. Type hints exist to prevent you breaking stuff but some "bad" things are allowed to happen if they aren't problematic (I'm looking at you int/float 👀)
  • Not all type checkers behave the same!! mypy and pyright show different behaviour in some edge cases...