Skip to content

← Modelling Data with Dataclasses step 25 of 28

Hard Framework

Type-Check a Custom Dataclass Decorator with dataclass_transform

Your team wants every model class to use the same dataclass options and to appear in a registry by name. So you write a decorator, @entity, that applies @dataclass(frozen=True, slots=True, kw_only=True) and registers the class. It works when the program runs, but mypy does not know that @entity makes a dataclass:

@entity
class Account:
    name: str


Account(name="alice")  # mypy: Unexpected keyword argument "name" for "Account"

typing.dataclass_transform fixes that. It marks a decorator as one that turns a class into a dataclass, so mypy checks the class as it checks @dataclass. It only informs the type checker. It does not write __init__, freeze fields or add slots when the program runs, so entity must still call dataclass itself.

Write entity_field

entity_field works like dataclasses.field, with one extra argument, secret. Keep its signature. It returns a field(...) with metadata={"secret": secret}, and it handles default and default_factory like this:

You pass entity_field does
Neither Returns a required field with the metadata
default only Passes the value on, even None, False, 0 or ""
default_factory only Passes the function on without calling it
Both Raises ValueError, as field does

Compare each argument with MISSING using is, because a false value such as 0 is still a default that the caller passed. secret is only a label in the metadata. It does not hide the value from repr, and the tests expect the normal repr.

Write entity

Put this declaration above entity:

@dataclass_transform(
    frozen_default=True,
    kw_only_default=True,
    field_specifiers=(entity_field,),
)
  • frozen_default=True tells mypy that the fields are read-only, so it rejects account.name = "x".
  • kw_only_default=True tells mypy that the constructor takes only keyword arguments.
  • field_specifiers=(entity_field,) tells mypy that a call to entity_field describes a field, the way field() does, so mypy reads its default argument. Without it, mypy treats the call itself as the default, and a required field looks optional.

In entity, apply dataclass(frozen=True, slots=True, kw_only=True) to cls. Register the class it returns in REGISTRY under its __name__, and return that same class. Do not register cls itself. With slots=True, dataclass builds a new class, and the registry must hold the class that creates the objects.

Keep Account and Session as they are. @entity runs when Python defines them, so REGISTRY is full before solve runs.

Write solve

solve(target, name) builds Account(name=name) for "Account" and Session(account=name) for "Session". Any other target raises ValueError. It returns a dict with these keys. Call the supplied helpers instead of returning the strings they produce.

Key Value
"registry" sorted(REGISTRY)
"registry_matches" REGISTRY[target] is type(obj)
"repr" repr(obj)
"frozen_error" _frozen_error(obj, "name") for an Account, or _frozen_error(obj, "account") for a Session
"positional_error" _positional_error(type(obj), name)
"secret_fields" The names of the fields whose metadata has a true "secret", in declaration order
"has_dict" hasattr(obj, "__dict__")

solve("Account", "alice") returns:

{
    "registry": ["Account", "Session"],
    "registry_matches": True,
    "repr": "Account(name='alice', token='', region='eu')",
    "frozen_error": "FrozenInstanceError",
    "positional_error": "TypeError",
    "secret_fields": ["token"],
    "has_dict": False,
}

A Session has ttl=60 and no secret fields.

Type checking

  • Keep the Any annotations on entity_field. Its result is the default for fields of different types, such as token: str and ttl: int. mypy accepts an Any result as the default for a field of any type.
  • Keep the generic signature entity[T](cls: type[T]) -> type[T], so a decorated class keeps its own type.
  • mypy does not check the dataclass_transform declaration against what entity does. If entity left out frozen=True, mypy would still reject assignments that the program then allowed. The checks in solve catch that.