pythonpython-typingmagic-methods

Python 3.10 - How to properly type hint __add__ in a dataclass that inherits Generic[TypeVar]?


I have this package called totemp that provides convertions between temperature scales. It's currently in development and I just don't know how to type hint the add magic method.

There are 8 dataclasses, Celsius, Fahrenheit, Delisle, Kelvin, Newton, Rankine, Réaumur and Romer.

Celsius class with dunder add and one of the convertion methods (@no_type_check to prevent mypy from not letting me commit).

from dataclasses import dataclass, field
from typing import Generic, TypeVar, no_type_check

TEMP = TypeVar('TEMP', int, float)  # TEMP must be int or float


@dataclass
class Celsius(Generic[TEMP]):
    """..."""

    __value: TEMP
    __symbol: str = field(compare=False, repr=False, default='ºC')

    def __str__(self) -> str:
        return f'{self.__value}{self.__symbol}'

    def __repr__(self) -> str:
        return f'Celsius({self.__value})'

    @no_type_check
    def __add__(self, other):
        match str(type(other)):
            case "<class 'totemp.temperature_types.Celsius'>":
                return Celsius(self.__value + other.value)
            case _:
                try:
                    other = other.to_celsius()
                    return Celsius(self.__value + other.value)
                except AttributeError as error:
                    print(
                        f'\033[31;1mAttributeError: {error}.\n    ' +
                        f"Cause: '{other}' is not a temperature scale.\033[m"
                    )

    def to_fahrenheit(self) -> 'Fahrenheit[TEMP]':
        """
        Returns a Fahrenheit object which contains the class attribute "value"
        with the result from the conversion typed the same as the attribute.

        Returns
        -------
        Fahrenheit[TEMP]
        """
        fahrenheit = type(self.__value)(self.__value * 9 / 5 + 32)
        return Fahrenheit(fahrenheit)

I tried creating an abstract class, but that didn't worked well, didn't even know what I was doing tbh. Tried read Python docs, but I had no response for what i'm trying to do, couldn't find any solution for myself. Thought about NewType, but after reading about in the docs, didn't understood it very well too.

Thought about doing something like:

GenericType = NewType('Temp', Union[Celsius, Fahrenheit, ...])

But ended up with more mypy errors that I couldn't deal with.

How do I properly type hint dunder add and other magic methods (such as gt, lt, e.g.)?


Solution

  • I got lost in more and more things I thought might be optimized with the code, so I ended up placing a merge request for your current dev branch.

    To quickly answer the basic idea relevant for this particular question:

    I don't think the class should be generic. I don't understand the purpose of that.

    A base class is definitely a good idea (considering the DRY principle). Making it abstract, you could enforce a convert_to method for all subclasses, which you could then utilize in your numeric operand dunder methods like __add__.

    Then typing would become fairly easy with bounded type variables:

    from abc import ABCMeta, abstractmethod
    from typing import Any, ClassVar, TypeVar
    
    
    T = TypeVar("T", bound="AbstractTemperature")
    
    
    class AbstractTemperature(metaclass=ABCMeta):
        def __init__(self, value: float) -> None:
            self._value = value
    
        @property
        def value(self) -> float:
            return self._value
    
        @abstractmethod
        def convert_to(self, temp_cls: type[T]) -> T:
            ...
    
        def __add__(self: T, other: Any) -> T:
            cls = self.__class__
            try:
                if isinstance(other, AbstractTemperature):
                    return cls(self._value + other.convert_to(cls).value)
                return cls(self._value + other)
            except TypeError:
                return NotImplemented