How to annotate that a function produces a dataclass?

Question

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)

Answer 1

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.

Answer 2

The problem is that mypy understands the metaclass decorator and its magic about __init__, but does not understand the dataclass function:

@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__)

As you can see, the __init__ method of Test2 does not accept any arguments.