Enforcing units on numbers using Python type hints

Question

Is there a way to use Python type hints as units? The type hint docs show some examples that suggest it might be possible using NewType, but also those examples show that addition of two values of the same "new type" do not give a result of the "new type" but rather the base type. Is there a way to enrich the type definition so that you can specify type hints that work like units (not insofar as they convert, but just so that you get a type warning when you get a different unit)? Something that would allow me to do this or similar:

Seconds = UnitType('Seconds', float)
Meters = UnitType('Meters', float)

time1 = Seconds(5)+ Seconds(8) # gives a value of type `Seconds`
bad_units1 = Seconds(1) + Meters(5) # gives a type hint error, but probably works at runtime
time2 = Seconds(1)*5 # equivalent to `Seconds(1*5)` 
# Multiplying units together of course get tricky, so I'm not concerned about that now.

I know runtime libraries for units exist, but my curiosity is if type hints in python are capable of handling some of that functionality.

Answer 1

You can do this by creating a type stub file, which defines the acceptable types for the __add__/__radd__ methods (which define the + operator) and __sub__/__rsub__ methods (which define the - operator). There are many more similar methods for other operators of course, but for the sake of brevity this example only uses those.

units.py

Here we define the units as simple aliases of int. This minimises the runtime cost, since we aren't actually creating a new class.

Seconds = int
Meters = int

units.pyi

This is a type stub file. It tells type checkers the types of everything defined in units.py, instead of having the types defined within the code there. Type checkers assume this is the source of truth, and don't raise errors when it differs from what is actually defined in units.py.

from typing import Generic, TypeVar

T = TypeVar("T")

class Unit(int, Generic[T]):
    def __add__(self, other: T) -> T: ...
    def __radd__(self, other: T) -> T: ...
    def __sub__(self, other: T) -> T: ...
    def __rsub__(self, other: T) -> T: ...
    def __mul__(self, other: int) -> T: ...
    def __rmul__(self, other: int) -> T: ...

class Seconds(Unit["Seconds"]): ...

class Meters(Unit["Meters"]): ...

Here we define Unit as a generic type inheriting from int, where adding/subtracting takes and returns values of type parameter T. Seconds and Meters are then defined as subclasses of Unit, with T equal to Seconds and Meters respectively.

This way, the type checker knows that adding/subtracting with Seconds takes and returns other values of type Seconds, and similarly for Meters.

Also, we define __mul__ and __rmul__ on Unit as taking a parameter of type int and returning T - so Seconds(1) * 5 should have type Seconds.

main.py

This is your code.

from units import Seconds, Meters

time1 = Seconds(5) + Seconds(8)
# time1 has type Seconds, yay!

bad_units1 = Seconds(1) + Meters(5)
# I get a type checking error:
# Operator "+" not supported for types "Meters" and "Seconds"
# Yay!

time2 = Seconds(1) * 5
# time2 has type Seconds, yay!

meter_seconds = Seconds(1) * Meters(5)
# This is valid because `Meters` is a subclass of `int` (as far
# as the type checker is concerned). meter_seconds ends up being
# type Seconds though - as you say, multiplying gets tricky.

---

Of course, all of this is just type checking. You can do what you like at run time, and the pyi file won't even be loaded.

Answer 2

The answer from @Artemis is excellent, but throws error when used with MyPy (@Artemis is using Pylance).

I made the following modifications to units.pyi (based on a suggestion from @Artemis) and it appears to be working well:

from typing import Generic, TypeVar, Union

T = TypeVar("T")

class Unit(Generic[T]):
    def __add__(self, other: Union[T, int]) -> T: ...
    def __radd__(self, other: Union[T, int]) -> T: ...
    def __sub__(self, other: Union[T, int]) -> T: ...
    def __rsub__(self, other: Union[T, int]) -> T: ...
    def __mul__(self, other: Union[T, int]) -> T: ...
    def __rmul__(self, other: Union[T, int]) -> T: ...

    def __init__(self, val: int) -> None: ...

class Seconds(Unit["Seconds"]): ...

class Meters(Unit["Meters"]): ...

The only hold-up is that you must create values using

v: Seconds = Seconds(1)

rather than:

v: Seconds = 1

Other than that, MyPy is able to catch operations using mixed types.