velai/serialization/JsonSerializable.py

import dataclasses
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Dict, TypeVar

from velai.serialization.JsonTypeSerializer import DefaultSerializer
from velai.serialization.serialization_utils import (
    resolve_type_hints,
    normalize_type_for_deserialization,
    normalize_type_for_serialization,
)

T = TypeVar("T")


class JsonSerializable(ABC):
    @abstractmethod
    def to_dict(self: T) -> Dict[str, Any]:
        raise NotImplementedError()

    @classmethod
    @abstractmethod
    def from_dict(cls: type[T], data: dict[str, Any]) -> T:
        raise NotImplementedError()

    @abstractmethod
    def update_from_dict(self: T, data: dict[str, Any]) -> None:
        raise NotImplementedError()


@dataclass
class DataclassJsonSerializable(JsonSerializable):
    def to_dict(self: T) -> Dict[str, Any]:
        result: Dict[str, Any] = {}
        type_hints = resolve_type_hints(type(self))

        for field in dataclasses.fields(self):
            value = getattr(self, field.name)
            hinted_type = type_hints[field.name]

            source_type = normalize_type_for_serialization(hinted_type)

            result[field.name] = DefaultSerializer.serialize(
                value,
                source_type=source_type,
            )

        return result

    @classmethod
    def from_dict(cls: type[T], data: Dict[str, Any]) -> T:
        parsed_data = cls._parse_fields(data)
        return cls(**parsed_data)

    def update_from_dict(self: T, data: dict[str, Any]) -> None:
        parsed_data = type(self)._parse_fields(data)
        for key, value in parsed_data.items():
            setattr(self, key, value)

    @classmethod
    def _parse_fields(cls, data: Dict[str, Any]) -> Dict[str, Any]:
        type_hints = resolve_type_hints(cls)
        parsed_data: Dict[str, Any] = {}

        for key, value in data.items():
            if key not in type_hints:
                continue

            if value is None:
                parsed_data[key] = None
                continue

            target_type = normalize_type_for_deserialization(type_hints[key])

            parsed_data[key] = DefaultSerializer.de_serialize(
                value,
                target_type=target_type,
            )

        return parsed_data