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