Search code examples
pythoninheritancepython-dataclassespython-class

How to inherit (from a parent class) dataclass field introspection functionality?


I have a parent dataclass, and various other classes will then extend this parent dataclass. Let's call these dataclasses DCs. In the example code below, see ParentDC, and an example ChildDC:

from dataclasses import dataclass, field, fields
from typing import Optional


@dataclass
class ParentDC:
    type_map: dict[str, type] = field(init=False, metadata={"support": True})
    primary_fields: set[str] = field(
        init=False, default_factory=set, metadata={"support": True}
    )
    secondary_fields: set[str] = field(
        init=False, default_factory=set, metadata={"support": True}
    )
    dtype_map: dict[str, type] = field(
        init=False, default_factory=dict, metadata={"support": True}
    )

    def __init_subclass__(cls, type_map: dict[str, type]) -> None:
        print(cls.__class__.__qualname__)
        cls.type_map = type_map

        cls.primary_fields = set()
        cls.secondary_fields = set()
        field_data = fields(cls)

        for fdat in field_data:
            if not fdat.metadata.get("support", False):
                if fdat.metadata.get("secondary", False):
                    cls.secondary_fields.add(fdat.name)
                else:
                    cls.primary_fields.add(fdat.name)

        cls.dtype_map = {
            k: type_map[v].dtype
            for k, v in cls.__annotations__.items()
            if k in cls.primary_fields.union(cls.secondary_fields)
        }


type_map = {
    "alpha": int,
    "beta": float,
}


@dataclass
class ChildDC(ParentDC, type_map=type_map):
    alpha: Optional[str] = field(
        default=None, kw_only=True, metadata={"secondary": True}
    )
    beta: str = field(kw_only=True)


print(f"{ChildDC.primary_fields=}")
print(f"{ChildDC.secondary_fields=}")
print(f"{ChildDC.dtype_map=}")

I want to create some "introspection functionality" common to all DCs. This introspection functionality rests at the class level: you should not need to create an instance of ChildDC to be able to access which fields are its primary fields, etc.

The code as it stands does not work:

type
ChildDC.primary_fields=set()
ChildDC.secondary_fields=set()
ChildDC.dtype_map={}

And I have some inkling of why: when __init_subclass__ is the wrong place for such introspection data to be generated and stored, because the parent does not have access to any of the child in __init_subclass__.

Where should this introspection information be generated and visualized?


Solution

  • One simple workaround is to perform the dataclass transformation in __init_subclass__ instead of decorating the subclass with the @ syntax so that the fields can be made available to the logics in your __init_subclass__:

    @dataclass
    class ParentDC:
        def __init_subclass__(cls, type_map: dict[str, type]) -> None:
            dataclass(cls)
            cls.type_map = type_map
            cls.primary_fields = set()
            cls.secondary_fields = set()
            field_data = fields(cls)
            ...
    
    class ChildDC(ParentDC, type_map=type_map):
        ...