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平分等于多少)

是否有一种方法在mymyy注释这个特定签名的函数?

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