Mixin for Nested Model support in SQLModel

sqlmodel_pydantic_json_mixin.py

def is_container_type(type_: Any) -> bool:
    return type_ in (list, set, dict, tuple)


def is_union_type(type_: Any) -> bool:
    return type_ is typing.Union or type_ is types.UnionType


def convert_field_value(annotation: Any, raw_value: Any) -> Any:
    origin = get_origin(annotation)
    args = get_args(annotation)

    if not origin:  # not a container type (e.g. int, untyped list, None, datetime)
        return raw_value

    if is_union_type(origin):
        for arg in args:
            converted = convert_field_value(arg, raw_value)
            if converted is not None:
                return converted
        return None

    if origin is list:
        return [convert_field_value(args[0], item) for item in raw_value]

    if origin is dict:
        return {
            key: convert_field_value(args[1], value) for key, value in raw_value.items()
        }

    if origin is tuple:
        if len(args) != len(raw_value):
            return raw_value
        return tuple(
            convert_field_value(arg, item) for arg, item in zip(args, raw_value)
        )

    if issubclass(annotation, PydanticBaseModel):
        model = annotation.model_construct()
        for field_name, field_info in annotation.model_fields.items():
            setattr(
                model,
                field_name,
                convert_field_value(field_info.annotation, raw_value.get(field_name)),
            )
        return model

    raise ValueError(f"Unsupported type {annotation}")


class PydanticJSONMixin:
    @reconstructor
    def init_on_load(self):
        for field_name, field_info in self.model_fields.items():
            raw_value = getattr(self, field_name)
            converted = convert_field_value(field_info.annotation, raw_value)
            setattr(self, field_name, converted)