如何为返回泛型实例与id属性映射字典的Python方法标注类型

实现方案 #

原有实现冗余的__new__协议方法定义、循环绑定的类型变量会干扰mypy的类型推导，简化协议约束即可满足需求：

from typing import Iterable, Protocol, TypeVar

# 定义id属性的类型变量
IdType = TypeVar("IdType")
# 定义实例类型变量，绑定到实现了HasId协议的类
ItemType = TypeVar("ItemType", bound="HasId[IdType]")

class HasId(Protocol[IdType]):
    @property
    def id(self) -> IdType: ...

def id_map(items: Iterable[ItemType]) -> dict[IdType, ItemType]:
    return {item.id: item for item in items}

验证示例 #

你可以用以下代码测试类型推导效果：

class User:
    def __init__(self, uid: int, name: str):
        self.id = uid
        self.name = name

users = [User(1, "张三"), User(2, "李四")]
# mypy 会自动推导 user_map 的类型为 dict[int, User]
user_map = id_map(users)

说明 #

1. 协议只需要定义你用到的id属性即可，不需要额外定义__new__等无关方法
2. 不需要额外的cast强转操作，mypy会通过结构子类型自动匹配符合协议的实例
3. 无特殊协变场景时不需要给TypeVar加covariant参数，避免增加不必要的复杂度

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