Pydantic Refactoring

python/lib/communication/dmod/communication/data_transmit_message.py

@@ -1,10 +1,23 @@
+from dmod.core.serializable import Serializable
+from pydantic import Extra
+from dmod.core.serializable_dict import SerializableDict
 from .message import AbstractInitRequest, MessageEventType, Response
-from typing import Dict, Optional, Union
-from numbers import Number
+from pydantic import Field
+from typing import ClassVar, Type, Union
 from uuid import UUID
 
 
-class DataTransmitMessage(AbstractInitRequest):
+class DataTransmitUUID(Serializable):
+    series_uuid: UUID = Field(description="A unique id for the collective series of transmission message this instance is a part of.")
+    """
+    The expectation is that a larger amount of data will be broken up into multiple messages in a series.
+    """
+
+    class Config:
+        field_serializers = {"series_uuid": lambda s: str(s)}
+
+
+class DataTransmitMessage(DataTransmitUUID, AbstractInitRequest):
     """
     Specialized message type for transmitting data.
 
@@ -18,64 +31,14 @@ class DataTransmitMessage(AbstractInitRequest):
     ::class:`str` object.  However, instances can be initialized using either ::class:`str` or ::class:`bytes` data.
     """
 
-    _KEY_SERIES_UUID = 'series_uuid'
+    event_type: ClassVar[MessageEventType] = MessageEventType.DATA_TRANSMISSION
 
-    event_type: MessageEventType = MessageEventType.DATA_TRANSMISSION
+    data: str = Field(description="The data carried by this message, in decoded string form.")
+    is_last: bool = Field(False, description="Whether this is the last data transmission message in this series.")
 
-    @classmethod
-    def factory_init_from_deserialized_json(cls, json_obj: dict) -> Optional['DataTransmitMessage']:
-        try:
-            return cls(data=json_obj['data'], series_uuid=UUID(json_obj[cls._KEY_SERIES_UUID]),
-                       is_last=json_obj['is_last'])
-        except Exception as e:
-            return None
 
-    def __init__(self, data: Union[str, bytes], series_uuid: UUID, is_last: bool = False, *args, **kwargs):
-        super(DataTransmitMessage, self).__init__(*args, **kwargs)
-        self._data: str = data if isinstance(data, str) else data.decode()
-        self._series_uuid = series_uuid
-        self._is_last: bool = is_last
-
-    @property
-    def data(self) -> str:
-        """
-        The data carried by this message, in decoded string form.
-
-        Returns
-        -------
-        str
-            The data carried by this message, in decoded string form.
-        """
-        return self._data
-
-    @property
-    def is_last(self) -> bool:
-        """
-        Whether this is the last data transmission message in this series.
-
-        Returns
-        -------
-        bool
-            Whether this is the last data transmission message in this series.
-        """
-        return self._is_last
-
-    @property
-    def series_uuid(self) -> UUID:
-        """
-        A unique id for the collective series of transmission message this instance is a part of.
-
-        The expectation is that a larger amount of data will be broken up into multiple messages in a series.
-
-        Returns
-        -------
-        UUID
-            A unique id for the collective series of transmission message this instance is a part of.
-        """
-        return self._series_uuid
-
-    def to_dict(self) -> Dict[str, Union[str, Number, dict, list]]:
-        return {'data': self.data, self._KEY_SERIES_UUID: str(self.series_uuid), 'is_last': self.is_last}
+class DataTransmitResponseBody(SerializableDict, DataTransmitUUID):
+    ...
 
 
 class DataTransmitResponse(Response):
@@ -86,38 +49,24 @@ class DataTransmitResponse(Response):
     series of which it is a part.
     """
 
-    response_to_type = DataTransmitMessage
-
-    _KEY_SERIES_UUID = response_to_type._KEY_SERIES_UUID
-
-    @classmethod
-    def factory_init_from_deserialized_json(cls, json_obj: dict) -> 'DataTransmitResponse':
-        """
-        Factory create a new instance of this type based on a JSON object dictionary deserialized from received JSON.
-
-        Parameters
-        ----------
-        json_obj
-
-        Returns
-        -------
-        response_obj : Response
-            A new object of this type instantiated from the deserialize JSON object dictionary, or none if the provided
-            parameter could not be used to instantiated a new object.
-        """
-        try:
-            return cls(success=json_obj['success'], reason=json_obj['reason'], message=json_obj['message'],
-                       series_uuid=json_obj['data'][cls._KEY_SERIES_UUID], data=json_obj['data'])
-        except Exception as e:
-            return None
-
-    def __init__(self, series_uuid: Union[str, UUID], *args, **kwargs):
-        if 'data' not in kwargs:
-            kwargs['data'] = dict()
-        kwargs['data'][self._KEY_SERIES_UUID] = str(series_uuid)
-        super(DataTransmitResponse, self).__init__(*args, **kwargs)
+    response_to_type: ClassVar[Type[AbstractInitRequest]] = DataTransmitMessage
+
+    data: DataTransmitResponseBody
+
+    # `series_uuid` required in prior version of code
+    def __init__(self, series_uuid: Union[str, UUID] = None, **kwargs):
+        # assume no need for backwards compatibility
+        if series_uuid is None:
+            super().__init__(**kwargs)
+            return
+
+        if "data" not in kwargs:
+            kwargs["data"] = dict()
+
+        kwargs["data"]["series_uuid"] = series_uuid
+        super().__init__(**kwargs)
 
     @property
     def series_uuid(self) -> UUID:
-        return UUID(self.data[self._KEY_SERIES_UUID])
+        return self.data.series_uuid