如何为数据类方法自动生成基于数据类成员的类型注解？

嘿，我之前也纠结过这个问题——给数据类的方法（比如用来更新属性的工具方法）自动生成和数据类字段完全匹配的类型注解，不想每次改了数据类字段还要手动同步方法的注解，重复劳动太烦人了！

先给你看我最初写的代码，就是普通的数据类加一个更新属性的方法：

from dataclasses import dataclass
import dataclasses
from typing import Self

@dataclass
class Dataclass:
    a: int
    b: str
    c: list[int]

    def update_attrs(self, **kwargs) -> Self:
        return self.__class__(**(dataclasses.asdict(self) | kwargs))

你看，这里的**kwargs根本没有具体的类型约束，类型检查器完全没法帮我们把关——比如我不小心传个d=123（不存在的字段），或者把a传成字符串，只有运行时才会报错，这就很容易留隐患。

后来我捣鼓出一个靠谱的办法：写一个装饰器，自动根据数据类的字段生成对应的类型注解，完全不用手动写重复内容！话不多说，直接上代码：

from dataclasses import dataclass, fields
from typing import Self, TypedDict, TypeVar, Partial, get_type_hints

# 定义泛型类型变量，用来约束数据类
T = TypeVar('T', bound=object)

def add_typed_update_method(cls: type[T]) -> type[T]:
    # 提取数据类的字段和对应的类型
    field_type_map = get_type_hints(cls)
    # 自动生成一个TypedDict，用来约束kwargs的类型，所有字段设为可选
    UpdateKwargs = TypedDict(
        f"{cls.__name__}UpdateKwargs",
        {name: field_type for name, field_type in field_type_map.items()},
        total=False
    )
    # 定义带正确类型注解的更新方法
    def update_attrs(self: T, **kwargs: Partial[UpdateKwargs]) -> Self:
        current_attrs = dataclasses.asdict(self)
        return self.__class__(**(current_attrs | kwargs))
    
    # 把方法绑定到数据类上
    setattr(cls, 'update_attrs', update_attrs)
    return cls

接下来用这个装饰器就很简单，只要在@dataclass下面加一层装饰器就行：

@add_typed_update_method
@dataclass
class Dataclass:
    a: int
    b: str
    c: list[int]

现在你用这个方法的时候，类型检查器就能准确识别允许的参数和对应的类型了，举个使用示例：

x = Dataclass(0, "1", [2])
# 类型检查器会认可这个调用，参数完全匹配字段类型
y = x.update_attrs(a=5, c=[3,4])
# 但如果这么写，类型检查器会直接报错：
# y = x.update_attrs(a="string", d=10)

以后你修改数据类的字段，比如加个d: float，注解会自动同步，完全不用手动改方法的类型约束，省了超多重复劳动！

这个办法的核心就是利用get_type_hints提取数据类的字段类型，动态生成对应的TypedDict作为kwargs的类型注解，完美实现了注解和数据类字段的同步。

内容来源于stack exchange