Mypy：注释函数，其中 return 类型取决于参数的类型

Question

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

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

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

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

    amount: int
    unit: MemoryUnit

现在我想为此 class 实现基本算法。加减乘法容易注释：

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

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

用Memory除以Memory得到float（两个存储量的比例是多少）
将Memory除以int得到Memory（如果n均分，Memory的数量是多少）

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

Mypy：注释函数，其中 return 类型取决于参数的类型

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

python

type-hinting

mypy

python-typing