velai/serialization/serialization_utils.py

import dataclasses
import types
from typing import Any, get_args, get_origin, get_type_hints, Union


def resolve_type_hints(cls: type) -> dict[str, Any]:
    try:
        hints = get_type_hints(cls, include_extras=True)
    except TypeError:
        hints = get_type_hints(cls)
    except Exception as e:
        raise TypeError(f"Failed to resolve type hints for {cls.__name__}: {e}") from e

    for field in dataclasses.fields(cls):
        if field.name not in hints:
            raise TypeError(f"{cls.__name__}.{field.name} has no type hint. All dataclass fields must be annotated.")

    return hints


def _is_union_origin(origin: Any) -> bool:
    return origin is Union or origin is types.UnionType


def normalize_type_for_deserialization(tp: Any) -> type:
    origin = get_origin(tp)

    if _is_union_origin(origin):
        args = [a for a in get_args(tp) if a is not type(None)]
        if len(args) == 1:
            return normalize_type_for_deserialization(args[0])
        raise TypeError(f"Unsupported Union type {tp}. Only Optional[T] or T | None are supported.")

    if origin is not None:
        return origin

    if isinstance(tp, type):
        return tp

    raise TypeError(f"Unsupported type annotation {tp!r}. Use a concrete runtime type.")


def normalize_type_for_serialization(tp: Any) -> type:
    origin = get_origin(tp)

    if _is_union_origin(origin):
        args = [a for a in get_args(tp) if a is not type(None)]
        if len(args) == 1:
            return normalize_type_for_serialization(args[0])
        raise TypeError(f"Unsupported Union type {tp}. Only Optional[T] or T | None are supported.")

    if isinstance(tp, type):
        return tp

    if origin is not None:
        return origin

    raise TypeError(f"Unsupported type annotation {tp!r} for serialization. Use a concrete runtime type.")