Typing and pint

Question

I'm using pint to use and convert units. I wanted to create classes which restricts the quantities only to "[time]" or "[length]" dimensions, so as a first approach I did the following:

from pint import Quantity, DimensionalityError

class Time(Quantity):
    def __new__(cls, v: str | Quantity) -> Quantity:
        obj = Quantity(v)
        if not obj.check("[time]"):
            raise DimensionalityError(v, "[time]")
        return obj

class Length(Quantity):
    def __new__(cls, v: str | Quantity) -> Quantity:
        obj = Quantity(v)
        if not obj.check("[length]"):
            raise DimensionalityError(v, "[length]")
        return obj

At runtime it works as expected, i.e: I can do the following:

1hour = Time("1h")    # Works ok, variable 1hour contains `<Quantity(1, 'hour')>`
bad = Time("1meter")  # As expected, raises pint.errors.DimensionalityError: Cannot convert from '1meter' to '[time]'
1meter = Length("1meter") # Ok
bad_again = Length("1h")  # Ok, raises DimensionalityError

However, from a typing perspective, something is wrong:

def myfunc(t: Time) -> str:
   return f"The duration is {t}"

print(myfunc(Time("1h")))    # Ok
print(myfunc(Length("1m")))  # Type error?

The second call to myfunc() is a type error, since I'm passing a Length instead of a Time. However mypy is happy with the code. So I have some questions:

1. Why doesn't mypy catches the error?
2. How to do it properly?

For 1. I guess that something fishy is happening in pint's Quantity implementation. I tried:

foo = Quantity("3 pounds")
reveal_type(foo)

and the revealed type is Any instead of Quantity which is very suspicious.

So I tried removing the base class Quantity from my Time and Length classes (i.e: they derive now from object instead of Quantity), and in this case, mypy correctly manages the typing errors.

But it fails again as soon as I try something like Length("60km")/Time("1h"). mypy complains that Length object does not implement the required method for performing that division (although the code works ok at runtime because, after all, Length and Time __new__() method is returning a Quantity object which does implement the arithmetical operations).

So, again, is there any workaround for making the idea work both at run-time and for mypy?

Answer 1

TL;DR: Pint has a wrapper function that can do that for you: @ureg.check('[time]') that goes right above your function's definition.

Discussion

In general, python does not enforce typing hints (see official documentation here). If you want to enforce such behavior, the best thing you can do is utilize the isinstance() native method (find details here). However, in your example, you are not even creating a new type. Both Time and Length classes will return a pint.util.Quantity object. Hence the following example would always raise an error, even for the Time defined objects:

def myfunc(t: Time) -> str:  # does not work!
if isinstance(t, Time):
    return f"The duration is {t}"
else:
    raise DimensionalityError(t, "[time]")

Solutions

1. A simple suggestion is to define a function that checks the quantity type:

   def qtype_check(obj: Quantity, q_type: str):
       qtype_str = f"[{q_type}]"
       if not obj.check(qtype_str):
           raise DimensionalityError(obj, qtype_str)
   

   And then simply call it in your method of choice:

   def myfunc(t: Quantity) -> str:
       qtype_check(t, 'time')
       return f"The duration is {t}"
   

2. A more elegant solution, if you need to keep your class definition, would be to additionally define your quantity classes as e.g. MyQuantity, with a class attribute _qtype that you change in each type definition. This way, you avoid writing the same error code snippet multiple times. Additionally, you can incorporate the qtype_check in the class, and invoke it through the type class of your choice (e.g.Time.qtype_check(value)), avoiding the hassle of having to remember how to define the class type string every time. Here are the definitions:

   from pint import Quantity, DimensionalityError
   
   
   class MyQuantity(Quantity):
       _qtype: str = "[]"
   
       def __new__(cls, v: str | Quantity) -> Quantity:
           obj = Quantity(v)
           cls.qtype_check(obj)
           return obj
   
       @classmethod
       def qtype_check(cls, obj: Quantity):
           if not obj.check(cls._qtype):
               raise DimensionalityError(obj, cls._qtype)
   
   
   class Time(MyQuantity):
       _qtype = "[time]"
   
   
   class Length(MyQuantity):
       _qtype = "[length]"
   
   
   def myfunc(t: Quantity) -> str:
       Time.qtype_check(t)
       return f"The duration is {t}"
   

3. Best of all is to simply use the pint-provided wrapper function, which can automatically check input units for you! Here is the documentation. In your case it would look like:

   @ureg.check('[time]')
   def myfunc(t: Quantity) -> str:
       return f"The duration is {t}"