如何注释函数生成数据类？

假设您想像这样包装 dataclass 装饰器：

from dataclasses import dataclass

def something_else(klass):
    return klass

def my_dataclass(klass):
    return something_else(dataclass(klass))

应该如何 my_dataclass 和/或 something_else 注释以指示返回类型是数据类？ 请参阅以下示例，了解内置 @dataclass 如何工作，但自定义 @my_dataclass 则不然：

@dataclass
class TestA:
    a: int
    b: str

TestA(0, "") # fine


@my_dataclass
class TestB:
    a: int
    b: str

TestB(0, "") # error: Too many arguments for "TestB" (from mypy)

要使的自定义装饰器 my_dataclass 正确地提示类型，需要使用 typing_extensions.dataclass_transform 装饰器。这个装饰器允许创建像 dataclass 一样的自定义装饰器，它可以正确地推断类型。

以下是如何修改的代码以实现此目的：

from dataclasses import dataclass
from typing_extensions import dataclass_transform

@dataclass_transform()
def something_else(klass):
    return klass

@dataclass_transform()
def my_dataclass(klass):
    return something_else(dataclass(klass))

@my_dataclass
class TestB:
    a: int
    b: str

TestB(0, "")  # 现在可以正常工作了

解释：

1. @dataclass_transform() : 这个装饰器告诉类型检查器，被装饰的函数（ something_else 和 my_dataclass ）正在转换一个类，并且转换后的类应该被视为数据类。

2. 两次使用 @dataclass_transform() : 需要在 something_else 和 my_dataclass 上都使用 @dataclass_transform() 。这是因为类型检查器需要了解这两个函数都参与了数据类的创建过程。

通过这些更改，mypy 和其他类型检查器将能够识别 my_dataclass 生成的类是数据类，并且不会再报告 "Too many arguments" 错误。