如何为带__getitem__、指定name/age键的Python Protocol写类型注解

正确类型注解实现方案 #

1. Protocol定义 #

你需要结合overload+Literal实现精准的__getitem__类型约束，完整代码如下：

from typing import Protocol, Literal, overload

class MyItemType(Protocol):
    # *若不需要支持点访问（item.name）、需兼容原生dict类型，可删除下方两个属性声明*
    name: str
    age: int

    @overload
    def __getitem__(self, key: Literal["name"]) -> str: ...
    @overload
    def __getitem__(self, key: Literal["age"]) -> int: ...
    # 兜底通用签名，可选添加
    def __getitem__(self, key: str) -> str | int: ...

2. 自定义类的__getitem__注解对应写法 #

你给出的MyItem类的__getitem__可以按如下方式注解：

from typing import Literal, overload, Union

class MyItem:
    def __init__(self, name:str, age: int):
        self.name = name
        self.age = age
    
    @overload
    def __getitem__(self, k: Literal["name"]) -> str: ...
    @overload
    def __getitem__(self, k: Literal["age"]) -> int: ...
    def __getitem__(self, k: str) -> Union[str, int]:
        if not hasattr(self, k):
            raise KeyError()
        return getattr(self, k)

原理解释 #

• 用Literal限定合法的键只能是name和age，类型检查器会自动拦截非法键的传入。
• 用overload声明不同键对应的返回值类型，可以保证item['name']自动推断为str，item['age']自动推断为int，类型精度更高。
• 该实现可以同时兼容pyspark.sql.Row、自定义符合要求的类，若删除属性声明也可以兼容带有name、age键的原生dict。
• 之前用TypedDict的方案无法满足需求，是因为TypedDict仅对字典结构生效，无法匹配实现了__getitem__的自定义类、Row等非dict类型。

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