typed.py

from abc import ABC
from copy import deepcopy
from dataclasses import dataclass, fields
from datetime import datetime
from enum import Enum
from typing import Union, get_args, get_origin
from uuid import UUID

nt = type(None)


def is_optional(t):
    args = get_args(t)
    return get_origin(t) is Union and len(args) > 1 and args[1] is nt


def parser(t, value):
    while is_optional(t):
        t = get_args(t)[0]
        if value is None:
            return value

    if get_origin(t) is dict:
        key_t, val_t = get_args(t)
        if not isinstance(value, dict):
            raise TypeError(f"value type is not dict instead of {type(value)}")
        return {
            parser(key_t, key_v): parser(val_t, val_v) for key_v, val_v in value.items()
        }

    if get_origin(t) is list:
        t = get_args(t)[0]
        if not isinstance(value, list):
            raise TypeError(f"value type is not list instead of {type(value)}")
        return [parser(t, val) for val in value]

    if issubclass(t, datetime):
        return datetime.fromtimestamp(int(value))
    if issubclass(t, _Common):
        return t.from_json_dict(value)
    if issubclass(t, Enum):
        return t(value)
    if issubclass(t, UUID):
        return UUID(value)

    return value


def encoder(t, value):
    while is_optional(t):
        t = get_args(t)[0]
        if value is None:
            return value

    if get_origin(t) is dict:
        key_t, val_t = get_args(t)
        if not isinstance(value, dict):
            raise TypeError(f"value type is not dict instead of {type(value)}")
        return {
            encoder(key_t, key_v): encoder(val_t, val_v)
            for key_v, val_v in value.items()
        }

    if get_origin(t) is list:
        t = get_args(t)[0]
        if not isinstance(value, list):
            raise TypeError(f"value type is not list instead of {type(value)}")
        return [encoder(t, val) for val in value]

    if issubclass(t, datetime):
        return int(value.timestamp())
    if issubclass(t, _Common):
        return value.to_json_dict()
    if issubclass(t, Enum):
        return value.value
    if issubclass(t, UUID):
        return str(value)

    return value


@dataclass
class _Common(ABC):
    @classmethod
    def from_json_dict(cls, obj):
        obj = deepcopy(obj)

        for field in fields(cls):
            if value := obj.get(field.name):
                obj[field.name] = parser(field.type, value)

        return cls(**obj)

    def to_json_dict(self):
        # Why not `asdict`? It will automatically
        # convert sub fields into dict, which
        # cause no match in below code
        obj = deepcopy(self.__dict__)

        for field in fields(self):
            if value := obj.get(field.name):
                obj[field.name] = encoder(field.type, value)

        return obj