Say you want to wrap the dataclass decorator like so:
from dataclasses import dataclass
def something_else(klass):
return klass
def my_dataclass(klass):
return something_else(dataclass(klass))
How should my_dataclass and/or something_else be annotated to indicate that the return type is a dataclass?
See the following example on how the builtin @dataclass works but a custom @my_dataclass does not:
@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)
Advertisement
Answer
There is no feasible way to do this prior to PEP 681.
A dataclass does not describe a type but a transformation. The actual effects of this cannot be expressed by Python’s type system – @dataclass is handled by a MyPy Plugin which inspects the code, not just the types. This is triggered on specific decorators without understanding their implementation.
dataclass_makers: Final = {
'dataclass',
'dataclasses.dataclass',
}
While it is possible to provide custom MyPy plugins, this is generally out of scope for most projects. PEP 681 (Python 3.11) adds a generic “this decorator behaves like @dataclass“-marker that can be used for all transformers from annotations to fields.
PEP 681 is available to earlier Python versions via typing_extensions.
Enforcing dataclasses
For a pure typing alternative, define your custom decorator to take a dataclass and modify it. A dataclass can be identified by its __dataclass_fields__ field.
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]:
...
This allows the type checker to understand and verify that a dataclass class is needed.
@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
Custom dataclass-like decorators
The PEP 681 dataclass_transform decorator is a marker for other decorators to show that they act “like” @dataclass. In order to match the behaviour of @dataclass, one has to use field_specifiers to indicate that fields are denoted the same way.
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))
It is possible for the custom dataclass decorator to take all keywords as @dataclass. dataclass_transform can be used to mark their respective defaults, even when not accepted as keywords by the decorator itself.