Update serde

src/py/flwr/common/serde.py

@@ -15,23 +15,32 @@
 """ProtoBuf serialization and deserialization."""
 
 
-from typing import Any, Dict, List, MutableMapping, OrderedDict, Type, TypeVar, cast
+from typing import (
+    Any,
+    Dict,
+    List,
+    MutableMapping,
+    OrderedDict,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from google.protobuf.message import Message as GrpcMessage
 
 # pylint: disable=E0611
+from flwr.proto.common_pb2 import BoolList, BytesList
+from flwr.proto.common_pb2 import ConfigsRecordValue as ProtoConfigsRecordValue
+from flwr.proto.common_pb2 import DoubleList, Sint64List, StringList
 from flwr.proto.error_pb2 import Error as ProtoError
 from flwr.proto.node_pb2 import Node
 from flwr.proto.recordset_pb2 import Array as ProtoArray
-from flwr.proto.recordset_pb2 import BoolList, BytesList
 from flwr.proto.recordset_pb2 import ConfigsRecord as ProtoConfigsRecord
-from flwr.proto.recordset_pb2 import ConfigsRecordValue as ProtoConfigsRecordValue
-from flwr.proto.recordset_pb2 import DoubleList
 from flwr.proto.recordset_pb2 import MetricsRecord as ProtoMetricsRecord
 from flwr.proto.recordset_pb2 import MetricsRecordValue as ProtoMetricsRecordValue
 from flwr.proto.recordset_pb2 import ParametersRecord as ProtoParametersRecord
 from flwr.proto.recordset_pb2 import RecordSet as ProtoRecordSet
-from flwr.proto.recordset_pb2 import Sint64List, StringList
 from flwr.proto.task_pb2 import Task, TaskIns, TaskRes
 from flwr.proto.transport_pb2 import (
     ClientMessage,
@@ -47,6 +56,7 @@
 from . import Array, ConfigsRecord, MetricsRecord, ParametersRecord, RecordSet, typing
 from .message import Error, Message, Metadata
 from .record.typeddict import TypedDict
+from .typing import ConfigsRecordValues
 
 #  === Parameters message ===
 
@@ -411,8 +421,8 @@ def _record_value_from_proto(value_proto: GrpcMessage) -> Any:
     return value
 
 
-def _record_value_dict_to_proto(
-    value_dict: TypedDict[str, Any],
+def record_value_dict_to_proto(
+    value_dict: Union[TypedDict[str, Any], Dict[str, ConfigsRecordValues]],
     allowed_types: List[type],
     value_proto_class: Type[T],
 ) -> Dict[str, T]:
@@ -431,7 +441,7 @@ def proto(_v: Any) -> T:
     return {k: proto(v) for k, v in value_dict.items()}
 
 
-def _record_value_dict_from_proto(
+def record_value_dict_from_proto(
     value_dict_proto: MutableMapping[str, Any]
 ) -> Dict[str, Any]:
     """Deserialize the record value dict from ProtoBuf."""
@@ -476,7 +486,7 @@ def parameters_record_from_proto(
 def metrics_record_to_proto(record: MetricsRecord) -> ProtoMetricsRecord:
     """Serialize MetricsRecord to ProtoBuf."""
     return ProtoMetricsRecord(
-        data=_record_value_dict_to_proto(record, [float, int], ProtoMetricsRecordValue)
+        data=record_value_dict_to_proto(record, [float, int], ProtoMetricsRecordValue)
     )
 
 
@@ -485,7 +495,7 @@ def metrics_record_from_proto(record_proto: ProtoMetricsRecord) -> MetricsRecord
     return MetricsRecord(
         metrics_dict=cast(
             Dict[str, typing.MetricsRecordValues],
-            _record_value_dict_from_proto(record_proto.data),
+            record_value_dict_from_proto(record_proto.data),
         ),
         keep_input=False,
     )
@@ -494,7 +504,7 @@ def metrics_record_from_proto(record_proto: ProtoMetricsRecord) -> MetricsRecord
 def configs_record_to_proto(record: ConfigsRecord) -> ProtoConfigsRecord:
     """Serialize ConfigsRecord to ProtoBuf."""
     return ProtoConfigsRecord(
-        data=_record_value_dict_to_proto(
+        data=record_value_dict_to_proto(
             record,
             [bool, int, float, str, bytes],
             ProtoConfigsRecordValue,
@@ -507,7 +517,7 @@ def configs_record_from_proto(record_proto: ProtoConfigsRecord) -> ConfigsRecord
     return ConfigsRecord(
         configs_dict=cast(
             Dict[str, typing.ConfigsRecordValues],
-            _record_value_dict_from_proto(record_proto.data),
+            record_value_dict_from_proto(record_proto.data),
         ),
         keep_input=False,
     )