Mypy:注释函数,其中 return 类型取决于参数的类型

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

我有一个 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: ...

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

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

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