Mypy：注释返回类型取决于参数类型的函数

Question

我有一个代表标准化内存量的类。

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

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

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: ...

不过我对除法有疑问。我看到除法的两个用例：

• 将

  Memory

  除以

  Memory

  并得到

  float

  （两个存储量之间的比率是多少）。
• 将

  Memory

  除以

  int

  ，得到

  Memory

  （如果将

  Memory

  均分的话，

  n

  的数量是多少）

mypy 中有没有办法用这个特定的签名来注释函数？

Answer 1

typing.overload

允许您注册一个函数的多个不同签名。用

@overload

修饰的函数在运行时会被忽略——它们只是为了类型检查器的利益——所以你可以将这些函数的主体留空。通常，您只需在这些函数的主体中放置文字省略号

...

、文档字符串或

pass

。另外，您需要确保至少有一个函数的具体实现可供在运行时使用。

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]"""