How can I make mypy work with a custom class decorator similar to dataclass?

Question

I need to generate code for classes in a very similar way dataclasses.dataclass does. In my first version, I wrote something like:

def typedrow(cls: type[_T]) -> type[_T]:
    cls_annotations = cls.__dict__.get("__annotations__", {})

    def __init__(s, *args, **kwargs) -> None:
        for field_name in cls_annotations.keys():
            s.__dict__[field_name] = kwargs.get(field_name)

    setattr(cls, "__init__", __init__)
    return cls

Although my code works, mypy is quite unhappy when I write:

@typedrow
class Person:
    name: Optional[str] = None

p = Person(name="joe")

It complains with error: Unexpected keyword argument "name" for "Person".

I noticed that dataclasses.dataclass does something completely different. Instead, it generates the code in text and calls _create_fn. Given that the generated code has all the type hints, it is kind obvious why mypy gets happy.

How does that work exactly?

I tried to create a simplified version of _create_fn which seems to be generated exactly the same code as dataclass, however, mypy is still unhappy.

Answer 1

mypy recognises the function dataclasses.dataclass by tracking how you import the name, then implementing custom type checking logic on the decorated class. This is all done in the mypy linting session, and has nothing to do with your Python runtime, so looking through the dataclasses.dataclass implementation is the wrong start.

PEP 681 - Data Class Transforms is dedicated for your use case. The following can be checked on mypy Playground:

from __future__ import annotations

import typing_extensions as t

if t.TYPE_CHECKING:
    _T = t.TypeVar("_T")

@t.dataclass_transform(eq_default=False, kw_only_default=True)
def typedrow(cls: type[_T], /) -> type[_T]:
    cls_annotations = cls.__dict__.get("__annotations__", {})

    def __init__(s: _T, *args: t.Any, **kwargs: t.Any) -> None:
        for field_name in cls_annotations.keys():
            s.__dict__[field_name] = kwargs.get(field_name)

    setattr(cls, "__init__", __init__)
    return cls

>>> @typedrow
... class Person:
...     name: str | None = None
...
>>> p = Person(name="joe")  # OK
>>> p = Person("joe")  # mypy: Too many positional arguments for "Person"  [misc]
>>> p = Person(name=1)  # mypy: Argument "name" to "Person" has incompatible type "int"; expected "str | None"  [arg-type]