如何让 mypy 使用类似于 dataclass 的自定义类装饰器？

Question

我需要以非常相似的方式为类生成代码

dataclasses.dataclass

。在我的第一个版本中，我写了类似的内容：

def typedrow(cls: type[_T]) -> type[_T]:
    cls_annotations = cls.__dict__.get("__annotations__", {})

    def __init__(s, *args, **kwargs) -> None:
        for field_name in cls_annotations.keys():
            s.__dict__[field_name] = kwargs.get(field_name)

    setattr(cls, "__init__", __init__)
    return cls

虽然我的代码可以工作，但是当我写的时候

mypy

很不高兴：

@typedrow
class Person:
    name: Optional[str] = None

p = Person(name="joe")

它抱怨

error: Unexpected keyword argument "name" for "Person"

。

我注意到

dataclasses.dataclass

做了完全不同的事情。相反，它会生成文本代码并调用

_create_fn

。鉴于生成的代码具有所有类型提示，很明显为什么

mypy

会感到高兴。

它是如何工作的到底？

我尝试创建一个简化版本的

_create_fn

，它似乎生成与

dataclass

完全相同的代码，但是，

mypy

仍然不高兴。

Answer 1

mypy 通过跟踪您如何导入名称来识别函数

dataclasses.dataclass

，然后在装饰类上实现自定义类型检查逻辑。这一切都是在 mypy linting 会话中完成的，与您的 Python 运行时无关，因此查看

dataclasses.dataclass

实现是错误的开始。

PEP 681 - 数据类转换 专用于您的用例。以下内容可以在 mypy Playground 上检查：

from __future__ import annotations

import typing_extensions as t

if t.TYPE_CHECKING:
    _T = t.TypeVar("_T")

@t.dataclass_transform(eq_default=False, kw_only_default=True)
def typedrow(cls: type[_T], /) -> type[_T]:
    cls_annotations = cls.__dict__.get("__annotations__", {})

    def __init__(s: _T, *args: t.Any, **kwargs: t.Any) -> None:
        for field_name in cls_annotations.keys():
            s.__dict__[field_name] = kwargs.get(field_name)

    setattr(cls, "__init__", __init__)
    return cls

>>> @typedrow
... class Person:
...     name: str | None = None
...
>>> p = Person(name="joe")  # OK
>>> p = Person("joe")  # mypy: Too many positional arguments for "Person"  [misc]
>>> p = Person(name=1)  # mypy: Argument "name" to "Person" has incompatible type "int"; expected "str | None"  [arg-type]