python - 如何注释一个函数产生一个数据类？

Question

假设你想像这样包装 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)

Answer 1

在 PEP 681 之前没有可行的方法来做到这一点.

dataclass 描述的不是类型而是转换。 Python 的类型系统无法表达其实际效果——@dataclass 由 MyPy Plugin 处理。它检查代码，而不仅仅是类型。这是 triggered on specific decorators不理解它们的实现。

dataclass_makers: Final = {
    'dataclass',
    'dataclasses.dataclass',
}

虽然可以提供自定义 MyPy 插件，但这通常超出了大多数项目的范围。 PEP 681(Python 3.11)添加了一个通用的“这个装饰器的行为类似于 @dataclass”-标记，可用于从注释到字段的所有转换器。

PEP 681 可通过 typing_extensions 用于早期的 Python 版本.

强制数据类

对于纯类型替代方案，定义您的自定义装饰器以接受数据类并对其进行修改。数据类可以通过其 __dataclass_fields__ 字段来标识。

from typing import Protocol, Any, TypeVar, Type
import dataclasses

class DataClass(Protocol):
    __dataclass_fields__: dict[str, Any]

DC = TypeVar("DC", bound=DataClass)

def my_dataclass(klass: Type[DC]) -> Type[DC]:
    ...

这允许类型检查器理解并验证需要dataclass类。

@my_dataclass
@dataclass
class TestB:
    a: int
    b: str

TestB(0, "")  # note: Revealed type is "so_test.TestB"

@my_dataclass
class TestC:  # error: Value of type variable "DC" of "my_dataclass" cannot be "TestC"
    a: int
    b: str

自定义类数据类装饰器

PEP 681 dataclass_transform 装饰器是其他 装饰器的标记，表明它们的行为“像”@dataclass。为了匹配 @dataclass 的行为，必须使用 field_specifiers 来指示字段以相同的方式表示。

from typing import dataclass_transform, TypeVar, Type
import dataclasses

T = TypeVar("T")

@dataclass_transform(
    field_specifiers=(dataclasses.Field, dataclasses.field),
)
def my_dataclass(klass: Type[T]) -> Type[T]:
    return something_else(dataclasses.dataclass(klass))

自定义数据类装饰器可以将所有关键字作为@dataclass。 dataclass_transform 可用于标记它们各自的默认值，即使装饰器本身不接受为关键字也是如此。