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

Question

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

from dataclasses import dataclass

def something_else(klass):
    return klass

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

How should 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 插件 用于检查代码，而不仅仅是类型。这是在特定装饰器上触发而不了解它们的实现。

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, ClassVar
from dataclasses import Field

class DataClass(Protocol):
    __dataclass_fields__: ClassVar[dict[str, Field[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 可用于标记它们各自的默认值，即使装饰器本身不接受作为关键字也是如此。

Answer 2

问题是 mypy 理解元类装饰器及其魔力
关于__init__，但不理解dataclass函数：

@dataclass
class Test:
    a: int


# Mypy: Revealed type is "def (self: Test, a: builtins.int)"
reveal_type(Test.__init__)


class Test2:
    a: int

#Mypy: Revealed type is "def (self: builtins.object)"
reveal_type(dataclass(Test2).__init__)

如您所见，Test2的__init__方法确实不接受任何论据。

Answer 3

当前接受的答案和评论描述了当前存在问题的原因。

存在以下大多数独立的问题：

1. 向类型检查器通知 __init__ 签名
2. 向类型检查器通知数据类。
3. 显式通过类型提示注释数据类。
4. 以通用方式执行 3。

1 & 2 由 dataclass 和 dataclass_transform 装饰器完成，或者更确切地说，它们是由类型检查器处理的。然而，对于人类来说，这在类型提示或注释中不可见，他们需要查看装饰器。
虽然它们不是注释，但它们解决了原始问题的问题。

2& 3 仅隐式可以通过 typeshed is_dataclass(obj) -> TypeIs[DataclassInstance] 形成条件交集类型。

4 对于人类来说只能通过 Annotated 或 TypeAliasType 完成（请参阅此答案的底部）。

---

下面我将回答有关如何显式注释此类函数或对象以及存在哪些问题的问题。

如果 Python 中的显式 Intersection 类型可以在 DataclassInstance 协议（请参阅本文结尾）和泛型之间形成子类：(type[T] ）->类型[DataclassInstance & T]。
目前最好的可能是通过类型检查器的条件隐式交集，它受T限制，因此非常有限：

def my_dataclass(klass : type[T]):
# -> type[<subclass of DataclassInstance and T.bound>]`
    obj = cast(Any, do_something(dataclass(klass)))
    assert is_dataclass(obj) and issubclass(obj, klass)
    return obj

类型检查器可以使用的上述类型保护是 python 的 typeshed 包提供的重载。
typeshed 是大多数类型检查器使用的官方库，它提供超出正常类型提示的类型提示带注释的模块的数量。 _typeshed 类型提示在运行时不可用，必须适当保护，例如 if TYPE_CHECKING，通过将类型提示封装在字符串，或使用 from __future__ import 注解。
虽然它看起来很简单并且您实际上可以做到这一点，但不鼓励在实际代码中包含_typeshed，它适用于存根文件（.pyi），即如果您使用它你应该为你的函数编写一个存根文件。

您的 IDE 类型检查器可能已经可以使用它，但如果需要，您可以通过以下方式安装它：

pip install types-dataclasses

注意：以下代码显示它可以显式完成，但它不是这是一个完美的解决方案，因为它不能以通用方式毫无问题地完成。

from typing import TYPE_CHECKING # not necessary in a stub file
if TYPE_CHECKING:                # not necessary in a stub file
   from _typeshed import DataclassInstance

# at best put this with @overload in a stub file!
@dataclass_transform(
    field_specifiers=(dataclasses.Field, dataclasses.field),
)
def my_dataclass(klass) -> type["DataclassInstance"]
   ...

@my_dataclass
class Foo:
  x : int
  y : list[int]

foo = Foo(1, [2]) # ok

@my_dataclass
class Baz(Foo):
  z : str

baz = Baz(1, [2], "b") # ok, because of dataclass_transform

if isinstance(bar, Foo): # deemed unnecessary
    print(bar)
else:
    assert_never(bar)  # ok

不起作用的是：

# Do not in a functional way

class Bar: ...

Bar_dc = my_dataclass(Bar)
assert is_dataclass(Bar_dc)    # error! Does not check for `type["DataclassInstance"]`
assert issubclass(Bar_dc, Bar) # Not True for type-checkers!

# ------

# Check for dataclass outside is the exact type is known.
if is_dataclass(foo):
    print("yes")      # reachable
else:
    assert_never(foo) # pyright: error, mypy: ok

if is_dataclass(Foo):
    print("yes")       # unreachable
else:
    assert_never(Foo) # reachable -> error

为什么最后的语句会失败？这是因为类型检查器不会将 Foo 视为 Foo & 的子类。数据类实例。目前这样的类型保护或特殊情况还不存在。对于 @dataclass ，这取决于类型检查器如何解释该类：

@dataclass
class Fox():
    a : int
    
if is_dataclass(Fox):
    reveal_type(Fox) # _pyright: unreachable | _mypy: Fox
else:
    reveal_type(Fox) # _pyright: Fox | _mypy: no-output

_typeshed.DataclassInstance

正如已经回答的，您可以使用 Protocol 和 _typeshed.DataclassInstance 实际上是一个。

from dataclasses import Field
from typing import ClassVar, Any, Protocol

# experimental
# Might not work as expected for pyright, see
#   https://github.com/python/typeshed/pull/9362
#   https://github.com/microsoft/pyright/issues/4339
class DataclassInstance(Protocol):
    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]

目前无法将对象显式注释为DataclassInstance并同时保留其原始类。

通用 TypeAliasType 注释

对类型检查器没有任何好处，但人类的通用注释提供了以下代码：

DataClass = TypeAliasType("DataClass", T, type_params=(T,))
"""
Describes a value that is a dataclass.

Note:
    This type-hint is similar to an annotation and does not
    provide any additional information to the type-checker, e.g. for `is_dataclass`.
"""

@dataclass_transform(...)
def make_dataclass(klass : type[T]) -> type[Dataclass[T]]: ...

# However without an additional cast the type-hint is not revealed in most situations.

foo = cast(DataClass[Foo], Foo(1, "yey")) # however reportUnnecessaryCast warning
# This cast displays the type of the instance as `DataClass[Foo]` in the tooltip instead of just `Foo`.

Answer 4

现有的答案似乎没有太多关于如何使用 dataclass_transform() 的信息，因此这里有一个实用指南。

---

@dataclass_transform 是一个装饰器，用于另一个装饰器函数、类或元类。这些装饰符号将成为“数据类工厂”。

基本用法

_{（游乐场：我的py，Pyright)}

# 1. Another decorator

@dataclass_transform()
def magic[T](cls: type[T]) -> type[T]: ...

@magic
class C:
    a: str
    b: int

reveal_type(C.__init__)  # (self: C, a: str, b: int) -> None

_{(游乐场：Mypy, Pyright)}

# 2. Another class

@dataclass_transform()
class Magic: ...

class C(Magic):
    a: str
    b: int

reveal_type(C.__init__)  # (self: C, a: str, b: int) -> None

_{(游乐场：Mypy, Pyright)}

# 3. A metaclass

@dataclass_transform()
class Magic(type): ...

class Base(metaclass=Magic):
    pass  # No fields

class C(metaclass=Magic):
    a: str
    b: int

class D(Base):
    a: str
    b: int

reveal_type(C.__init__)  # (self: C, a: str, b: int) -> None
reveal_type(D.__init__)  # (self: D, a: str, b: int) -> None

@dataclass_transform 的参数

根据工厂的行为方式，您可以提供 相应的关键字参数 @dataclass_transform：

_{（游乐场：Mypy, Pyright)}

# This means, if the `frozen` keyword-only argument is not provided,
# the decorated symbol will produce "frozen" instances by default.
@dataclass_transform(frozen_default = True)
def magic[T](*, frozen: bool = True) -> Callable[[type[T]], type[T]]: ...
#               ^^^^^^^^^^^^^^^^^^^ Not necessary if only frozen instances are desired.


@magic()  # or @magic(frozen = True)
class C:
    a: str
    b: int

c = C(a = 'lorem', b = 42)
c.a = ''  # error: Cannot modify attribute


@magic(frozen = False)
class D:
    a: str
    b: int

d = D(a = 'lorem', b = 42)
d.a = ''  # fine

字段和字段说明符

@dataclass_transform 还有一个名为 field_specifiers 的参数。赋予此参数的参数必须是可调用元组，每个可调用元组都可用于描述分配给它的字段的行为。

下面是一个如何使用 dataclasses.field 的小示例：

_{(playgrounds: Mypy, Pyright)}

@dataclass
class C:
    # Defaults to 0 if not given
    a: int = field(default = 0)
    # Calls the lambda to retrieve the default value if not given
    b: list[int] = field(default_factory = lambda: [1, 2, 3])

c = C()
reveal_type(c.a)  # int / 0
reveal_type(c.b)  # list[int] / [1, 2, 3]

以下是如何使用自定义装饰器和字段说明符执行相同操作：

_{(playgrounds: Mypy, Pyright)}

# Holder for field-related data.
# This should be read and handled by @magic internally.
@dataclass(kw_only = True)
class MagicField:
    default: Any = None
    default_factory: Callable[[], Any] | None = None

# Field specifier
def magic_field(
    *,
    default: Any = None,
    default_factory: Callable[[], Any] | None = None
) -> Any: ...
#    ^^^ This is important.
# At runtime, return `MagicField` so that @magic can read it.

# Can also be defined as multiple overloads
# (`default` and `default_factory` are mutually exclusive).
@overload
def magic_field[T](*, default: T | None = None) -> T: ...

@overload
def magic_field[T](*, default_factory: Callable[[], T] | None = None) -> T: ...

def magic_field[...](...) -> ...: ...  # Omitted here for brevity

#                                       vvvvvvvvvvvvvv
@dataclass_transform(field_specifiers = (magic_field,))
def magic(...) -> ...: ...  # Omitted here for brevity

@magic
class C:
    #        vvvvvv Prefix added, but nothing else
    a: int = magic_field(default = 0)
    b: list[int] = magic_field(default_factory = lambda: [1, 2, 3])

c = C()
reveal_type(c.a)  # int / 0
reveal_type(c.b)  # list[int] / [1, 2, 3]

可以找到字段说明符参数列表 规范中。

一个小警告：

_{(游乐场：Mypy, Pyright)}

@magic
class D:
    # Technically, parameters to field specifiers don't need
    # to be keyword-only. However, the corresponding arguments
    # must be passed as keyword arguments to be recognized.
    e: int = magic_field(0)  # silently ignored

d = D()  # error: Missing argument (0 is not recognized as default value)