refactor: separate message serialization and encoding

Author: David Perl

bec_lib/bec_lib/bec_serializable.py

@@ -0,0 +1,23 @@
+from pydantic import BaseModel, ConfigDict, computed_field
+
+
+class BecCodecInfo(BaseModel):
+    type_name: str
+
+
+class BECSerializable(BaseModel):
+    """A base class for serializable BEC objects, especially BEC messages.
+    Fields in subclasses which use non-primitive types must be in structured,
+    type-hinted objects, and their encoders and JSON schema should be defined in
+    this class."""
+
+    model_config = ConfigDict(
+        json_schema_serialization_defaults_required=True,
+        arbitrary_types_allowed=True,
+        extra="forbid",
+    )
+
+    @computed_field()
+    @property
+    def bec_codec(self) -> BecCodecInfo:
+        return BecCodecInfo(type_name=self.__class__.__name__)

bec_lib/bec_lib/devicemanager.py

@@ -667,9 +667,18 @@ def _get_redis_device_config(self) -> list:
 
     def _add_multiple_devices_with_log(self, devices: Iterable[tuple[dict, DeviceInfoMessage]]):
         logs = (self._add_device(*conf_msg) for conf_msg in devices if conf_msg is not None)
-        logger.info(f"Adding new devices:\n" + ", ".join(f"{name}: {t}" for name, t in logs))  # type: ignore # filtered
+        if set(logs) == {None}:
+            logger.warning("No devices added!")
+            return
+        logger.info(
+            f"Adding new devices:\n"
+            + ", ".join(f"{log[0]}: {log[1]}" for log in logs if log is not None)
+        )
 
     def _add_device(self, dev: dict, msg: DeviceInfoMessage) -> tuple[str, str] | None:
+        if msg is None:
+            logger.error(f"No device info in Redis for: {dev}")
+            return None
         name = msg.content["device"]
         info = msg.content["info"]
 

bec_lib/bec_lib/messages.py

@@ -9,7 +9,7 @@
 from importlib.metadata import PackageNotFoundError
 from importlib.metadata import version as importlib_version
 from types import NoneType
-from typing import Annotated, Any, ClassVar, Literal, Self, Union
+from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, Union
 from uuid import uuid4
 
 import msgpack
@@ -26,6 +26,7 @@
 )
 from typing_extensions import TypeAliasType
 
+from bec_lib.bec_serializable import BECSerializable
 from bec_lib.metadata_schema import get_metadata_schema_for_scan
 from bec_lib.one_way_registry import OneWaySerializationRegistry
 
@@ -113,7 +114,7 @@ class BECStatus(Enum):
     ERROR = -1
 
 
-class BECMessage(BaseModel):
+class BECMessage(BECSerializable):
     """Base Model class for BEC Messages
 
     Args:
@@ -122,7 +123,6 @@ class BECMessage(BaseModel):
 
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
     msg_type: ClassVar[str]
     metadata: JsonableDict = Field(default_factory=dict)
 
@@ -1307,16 +1307,19 @@ class DAPResponseMessage(BECMessage):
     dap_request: BECMessage | None = Field(default=None)
 
 
+MessageType = TypeVar("MessageType", bound=BECMessage)
+
+
 class AvailableResourceMessage(BECMessage):
     """Message for available resources such as scans, data processing plugins etc
 
     Args:
-        resource (dict, list[dict], BECMessage, list[BECMessage]): Resource description
+        resource (dict, list[dict], BECMessage, list[BECMessage]): Resource description - may contain only one type of BECMessage
         metadata (dict, optional): Metadata. Defaults to None.
     """
 
     msg_type: ClassVar[str] = "available_resource_message"
-    resource: JsonableDict | list[JsonableDict] | BECMessage | list[BECMessage]
+    resource: JsonableDict | list[JsonableDict] | MessageType | list[MessageType]
 
 
 class ProgressMessage(BECMessage):

bec_lib/bec_lib/serialization.py

@@ -36,6 +36,8 @@ class BECMessagePack(SerializationRegistry):
 
     def dumps(self, obj):
         """Pack object `obj` and return packed bytes."""
+        if isinstance(obj, BECMessage):
+            obj = obj.model_dump(mode="python", fallback=self.encode)
         return msgpack_module.packb(obj, default=self.encode)
 
     def loads(self, raw_bytes):

bec_lib/bec_lib/serialization_registry.py

@@ -4,6 +4,7 @@
 from typing import Callable, Type
 
 from bec_lib import codecs as bec_codecs
+from bec_lib import messages
 from bec_lib.logger import bec_logger
 
 logger = bec_logger.logger
@@ -18,7 +19,6 @@ def __init__(self):
         self._registry: dict[str, tuple[Type, Callable, Callable]] = {}
         self._legacy_codecs = []  # can be removed in future versions, see issue #516
 
-        self.register_codec(bec_codecs.BECMessageEncoder)
         self.register_codec(bec_codecs.EndpointInfoEncoder)
         self.register_codec(bec_codecs.SetEncoder)
         self.register_codec(bec_codecs.BECTypeEncoder)
@@ -97,6 +97,11 @@ def encode(self, obj):
 
     def decode(self, data):
         """Decode an object using the registered codec."""
+        if isinstance(data, dict) and "bec_codec" in data:
+            codec_info = data.pop("bec_codec")
+            msg_cls = messages.__dict__.get(codec_info.get("type_name"))
+            if msg_cls is not None:
+                return msg_cls.model_validate(data)
         if not isinstance(data, dict) or "__bec_codec__" not in data:
             return data
         codec_info = data["__bec_codec__"]

bec_lib/tests/test_bec_messages.py

@@ -17,7 +17,7 @@ def test_bec_message_msgpack_serialization_version(version):
         assert "Unsupported BECMessage version" in str(exception.value)
     else:
         res = MsgpackSerialization.dumps(msg)
-        res_expected = b"\x81\xad__bec_codec__\x83\xacencoder_name\xaaBECMessage\xa9type_name\xb8DeviceInstructionMessage\xa4data\x84\xa8metadata\x81\xa3RID\xa41234\xa6device\xa4samx\xa6action\xa3set\xa9parameter\x81\xa3set\xcb?\xe0\x00\x00\x00\x00\x00\x00"
+        res_expected = b"\x85\xa8metadata\x81\xa3RID\xa41234\xa6device\xa4samx\xa6action\xa3set\xa9parameter\x81\xa3set\xcb?\xe0\x00\x00\x00\x00\x00\x00\xa9bec_codec\x81\xa9type_name\xb8DeviceInstructionMessage"
         assert res == res_expected
         res_loaded = MsgpackSerialization.loads(res)
         assert res_loaded == msg
@@ -682,3 +682,14 @@ def test_message_with_np_array_in_dict():
     arr = np.zeros(5)
     msg = messages.ScanMessage(point_id=0, scan_id="", data={"device": {"value": arr}}, metadata={})
     assert isinstance(msg.data["device"]["value"], np.ndarray)
+
+
+def test_message_service_config():
+    msg = messages.MessagingServiceConfig(
+        metadata={}, service_name="signal", scopes=["*"], enabled=True
+    )
+    dump = msg.model_dump(mode="python")
+    assert dump["service_name"] == "signal"
+    resource_msg = messages.AvailableResourceMessage(resource=[msg])
+    resource_msg_dump = resource_msg.model_dump(mode="python")
+    assert resource_msg_dump["resource"][0]["service_name"] == "signal"

bec_server/bec_server/procedures/manager.py

@@ -50,6 +50,7 @@ def _log_on_end(future: Future):
 
 def _resolve_dict(msg: dict[str, Any] | _T, MsgType: type[_T]) -> _T:
     if isinstance(msg, dict):
+        msg.pop("bec_codec", None)
         return MsgType.model_validate(msg)
     return msg