Mypy: Annotate function where return type depends on the type of an argument

user_271789 0
Viewed 440

I have a class that represents a normalized amount of memory.

class MemoryUnit(enum.Enum):
    """Units of memory."""

    GB = 'GB'
    TB = 'TB'
    PB = 'PB'

class Memory(BaseModel):
    """Normalized amount of memory."""

    amount: int
    unit: MemoryUnit

Now I want to implement basic arithmetics for this class. Addition, substraction and multiplication are easy to annotate:

def __add__(self, other: Memory) -> Memory: ...
def __sub__(self, other: Memory) -> Memory: ...
def __mul__(self, other: int) -> Memory: ...

I have a problem with division though. I see two use-cases for division:

  • divide Memory by Memory and get a float (what is the ratio between two amounts of storage).
  • divide Memory by int and get Memory (what is the amount of Memory if divided evenly among n)

Is there a way in mypy to annotate a function with this specific signature?

python-typing mypy type-hinting python
1 Answers

typing.overload allows you to register multiple different signatures of a function. Functions decorated with @overload are ignored at runtime — they are just for the type-checker's benefit — so you can leave the body of these functions empty. Typically, you'll just put a literal ellipsis ..., a docstring, or pass in the body of these functions. Also, you'll need to make sure there is at least one concrete implementation of the function for use at runtime.

from typing import overload, Union

class Memory(BaseModel):
    """Normalized amount of memory."""

    amount: int
    unit: MemoryUnit

    def __add__(self, other: Memory) -> Memory: ...
    def __sub__(self, other: Memory) -> Memory: ...
    def __mul__(self, other: int) -> Memory: ...

    @overload
    def __div__(self, other: Memory) -> float:
        """Signature of `Memory.__div__` if an instance of `Memory`
        is divided by another instance of `Memory`
        """

    @overload
    def __div__(self, other: int) -> Memory:
        """Signature of `Memory.__div__` if an instance of `Memory`
        is divided by an `int`
        """

    def __div__(self, other: Union[int, Memory]) -> Union[Memory, float]:
        """[Your actual runtime implementation goes here]"""
user_139900160
Related