software/src/protos/serialization.py

1"""Serialization utilities for Protocol Buffer-like messages.23This module provides msgpack-based serialization for proto dataclasses,4with custom handling for datetime objects and hash computation for5reproducibility tracking.6"""78from __future__ import annotations910import hashlib11from dataclasses import asdict, is_dataclass12from datetime import datetime, timezone13from enum import IntEnum14from typing import Any, TypeVar1516import msgpack1718T = TypeVar("T")192021def _encode_value(obj: Any) -> Any:22    """Recursively encode a value for msgpack serialization.2324    Args:25        obj: Value to encode.2627    Returns:28        Msgpack-compatible value.29    """30    if isinstance(obj, datetime):31        return {"__datetime__": True, "iso": obj.isoformat()}32    elif isinstance(obj, IntEnum):33        return int(obj)34    elif is_dataclass(obj) and not isinstance(obj, type):35        return {k: _encode_value(v) for k, v in asdict(obj).items()}36    elif isinstance(obj, dict):37        return {k: _encode_value(v) for k, v in obj.items()}38    elif isinstance(obj, (list, tuple)):39        return [_encode_value(item) for item in obj]40    else:41        return obj424344def _decode_value(obj: Any, expected_type: type | None = None) -> Any:45    """Recursively decode a value from msgpack deserialization.4647    Args:48        obj: Value to decode.49        expected_type: Optional expected type for reconstruction.5051    Returns:52        Decoded value.53    """54    if isinstance(obj, dict):55        if obj.get("__datetime__"):56            iso_str = obj["iso"]57            # Handle both aware and naive datetimes58            dt = datetime.fromisoformat(iso_str)59            # Ensure timezone-aware60            if dt.tzinfo is None:61                dt = dt.replace(tzinfo=timezone.utc)62            return dt63        return {k: _decode_value(v) for k, v in obj.items()}64    elif isinstance(obj, list):65        return [_decode_value(item) for item in obj]66    else:67        return obj686970def serialize_proto(proto: Any) -> bytes:71    """Serialize a proto dataclass to msgpack bytes.7273    Args:74        proto: A dataclass instance to serialize.7576    Returns:77        Msgpack-encoded bytes.7879    Raises:80        TypeError: If proto is not a dataclass instance.81    """82    if not is_dataclass(proto) or isinstance(proto, type):83        raise TypeError(f"Expected dataclass instance, got {type(proto).__name__}")8485    encoded = _encode_value(proto)86    return msgpack.packb(encoded, use_bin_type=True)878889def deserialize_proto(data: bytes, proto_class: type[T]) -> T:90    """Deserialize msgpack bytes to a proto dataclass.9192    Args:93        data: Msgpack-encoded bytes.94        proto_class: The dataclass type to instantiate.9596    Returns:97        Instance of proto_class populated with decoded data.9899    Raises:100        TypeError: If proto_class is not a dataclass.101        ValueError: If data cannot be deserialized.102    """103    if not is_dataclass(proto_class):104        raise TypeError(f"Expected dataclass type, got {proto_class}")105106    try:107        decoded = msgpack.unpackb(data, raw=False)108    except Exception as e:109        raise ValueError(f"Failed to decode msgpack data: {e}") from e110111    # Recursively decode values112    decoded = _decode_value(decoded)113114    # Reconstruct the dataclass115    return _reconstruct_dataclass(decoded, proto_class)116117118def _reconstruct_dataclass(data: dict, proto_class: type[T]) -> T:119    """Reconstruct a dataclass from decoded data.120121    Args:122        data: Decoded dictionary data.123        proto_class: The dataclass type to instantiate.124125    Returns:126        Instance of proto_class.127    """128    import dataclasses129    import types130    import typing131132    if not isinstance(data, dict):133        raise ValueError(f"Expected dict, got {type(data).__name__}")134135    # Get field types for nested reconstruction, with type hints resolved136    type_hints = typing.get_type_hints(proto_class)137    fields = {f.name: f for f in dataclasses.fields(proto_class)}138    kwargs = {}139140    for name, value in data.items():141        if name not in fields:142            continue143144        # Use resolved type hints instead of field.type to handle forward refs145        field_type = type_hints.get(name, fields[name].type)146147        # Handle string type annotations that couldn't be resolved148        if isinstance(field_type, str):149            kwargs[name] = value150            continue151152        # Handle Optional types and Union types (X | None)153        origin = typing.get_origin(field_type)154        args = typing.get_args(field_type)155156        # Check for Union types (typing.Union or types.UnionType for X | Y syntax)157        is_union = origin is typing.Union or isinstance(field_type, types.UnionType)158        if is_union:159            # Get the actual type from Optional[X] or X | None160            # A Union always has args, and at least one is not NoneType161            # (`None | None` collapses to NoneType, which is not a Union).162            if args:  # pragma: no branch163                # Filter out NoneType164                non_none_types = [t for t in args if t is not type(None)]165                if non_none_types:  # pragma: no branch166                    field_type = non_none_types[0]167                    # Re-evaluate origin for the actual type168                    origin = typing.get_origin(field_type)169                    args = typing.get_args(field_type)170171        # Recursively reconstruct nested dataclasses172        if value is not None and is_dataclass(field_type) and isinstance(value, dict):173            value = _reconstruct_dataclass(value, field_type)174        elif origin is list and isinstance(value, list) and value:175            # Check if it's a list of dataclasses176            if args:177                list_item_type = args[0]178                if is_dataclass(list_item_type):179                    value = [180                        _reconstruct_dataclass(item, list_item_type)181                        if isinstance(item, dict)182                        else item183                        for item in value184                    ]185186        kwargs[name] = value187188    return proto_class(**kwargs)189190191def compute_proto_hash(proto: Any) -> str:192    """Compute a SHA256 hash of a serialized proto.193194    The hash is computed from the msgpack-serialized bytes, providing195    a deterministic identifier for the proto contents.196197    Args:198        proto: A dataclass instance to hash.199200    Returns:201        Hex-encoded SHA256 hash string.202203    Raises:204        TypeError: If proto is not a dataclass instance.205    """206    serialized = serialize_proto(proto)207    return hashlib.sha256(serialized).hexdigest()208