如何为Python对象unpacking添加type hints使Mypy正确推导类型？

问题代码与现象 #

现有代码尝试为自定义类实现解包的类型提示，但直接解包对象时Mypy无法推断变量类型：

from typing import Tuple


class A:
    def __init__(self, x: int, y: str):
        self.x = x
        self.y = y

    def astuple(self) -> Tuple[int, str]:
        return self.x, self.y

    # 需要为__iter__标注返回类型
    def __iter__(self):
        return iter(self.astuple())


a = A(1, "a")
# 无法推断x和y的类型
x, y = a
reveal_type(x)
reveal_type(y)
# 能正确推断p和q的类型为int和str
p, q = a.astuple()
reveal_type(p)
reveal_type(q)

运行mypy后的输出：

$ mypy unpack_object.py
unpack_object.py:20: note: Revealed type is "Any"
unpack_object.py:21: note: Revealed type is "Any"
unpack_object.py:24: note: Revealed type is "builtins.int"
unpack_object.py:25: note: Revealed type is "builtins.str"
Success: no issues found in 1 source file

直接解包对象得到的x、y被推断为Any，但解包astuple()返回的元组时，p、q的类型能被正确识别。

解决方案 #

要让Mypy正确推断对象解包后的类型，需让类满足类型检查器的序列识别规则，推荐以下两种方案：

方案一：实现Sequence协议 #

通过定义__len__和__getitem__方法，让Mypy将类视为固定长度的序列，从而精准推断解包后每个变量的类型：

from typing import Tuple, Sequence, Iterator


class A(Sequence[int | str]):
    def __init__(self, x: int, y: str):
        self.x = x
        self.y = y

    def astuple(self) -> Tuple[int, str]:
        return self.x, self.y

    def __iter__(self) -> Iterator[int | str]:
        return iter(self.astuple())
    
    def __len__(self) -> int:
        return 2
    
    def __getitem__(self, index: int) -> int | str:
        if index == 0:
            return self.x
        elif index == 1:
            return self.y
        raise IndexError("Index out of range")


a = A(1, "a")
x, y = a
reveal_type(x)  # 推断为int
reveal_type(y)  # 推断为str
p, q = a.astuple()
reveal_type(p)
reveal_type(q)

方案二：精确标注__iter__并结合类型断言 #

若不想实现Sequence协议，可给__iter__标注返回迭代器的元素类型，再通过类型断言指定解包变量的精确类型：

from typing import Tuple, Iterator, Union


class A:
    def __init__(self, x: int, y: str):
        self.x = x
        self.y = y

    def astuple(self) -> Tuple[int, str]:
        return self.x, self.y

    def __iter__(self) -> Iterator[Union[int, str]]:
        return iter(self.astuple())


a = A(1, "a")
# 直接解包会得到Union[int, str]类型
x, y = a
reveal_type(x)  # 推断为Union[int, str]
reveal_type(y)  # 推断为Union[int, str]

# 手动指定精确类型
x: int = a[0]
y: str = a[1]

验证结果 #

使用方案一运行mypy后，输出如下：

$ mypy unpack_object.py
unpack_object.py:25: note: Revealed type is "builtins.int"
unpack_object.py:26: note: Revealed type is "builtins.str"
unpack_object.py:28: note: Revealed type is "builtins.int"
unpack_object.py:29: note: Revealed type is "builtins.str"
Success: no issues found in 1 source file

内容的提问来源于stack exchange，提问作者Aravind Nujella