Refactor serialization module to simplify type conversion and enhance Pydantic integration

grongierisc · grongierisc · commit de27d318ae2c · 2025-01-28T15:21:13.000+01:00
diff --git a/src/iop/_serialization.py b/src/iop/_serialization.py
@@ -1,173 +1,52 @@
 from __future__ import annotations
-import base64
 import codecs
-import datetime
-import decimal
 import importlib
-import json
+import inspect
 import pickle
-import uuid
+import json
+from dataclasses import asdict, is_dataclass
 from typing import Any, Dict, Type
 
-from dacite import Config, from_dict
 import iris
+from pydantic import BaseModel, TypeAdapter
 
 from iop._message import _PydanticPickleMessage
 from iop._utils import _Utils
-from pydantic import BaseModel
-
-# Constants
-DATETIME_FORMAT_LENGTH = 23
-TIME_FORMAT_LENGTH = 12
-TYPE_SEPARATOR = ':'
-SUPPORTED_TYPES = {
-    'datetime', 'date', 'time', 'dataframe', 
-    'decimal', 'uuid', 'bytes'
-}
 
 class SerializationError(Exception):
-    """Base exception for serialization errors."""
+    """Exception raised for serialization errors."""
     pass
 
-class TypeConverter:
-    """Handles type conversion for special data types."""
-    
-    @staticmethod
-    def convert_to_string(typ: str, obj: Any) -> str:
-        if typ == 'dataframe':
-            return obj.to_json(orient="table")
-        elif typ == 'datetime':
-            return TypeConverter._format_datetime(obj)
-        elif typ == 'date':
-            return obj.isoformat()
-        elif typ == 'time':
-            return TypeConverter._format_time(obj)
-        elif typ == 'bytes':
-            return base64.b64encode(obj).decode("UTF-8")
-        return str(obj)
-
-    @staticmethod
-    def convert_from_string(typ: str, val: str) -> Any:
-        try:
-            if typ == 'datetime':
-                return datetime.datetime.fromisoformat(val)
-            elif typ == 'date':
-                return datetime.date.fromisoformat(val)
-            elif typ == 'time':
-                return datetime.time.fromisoformat(val)
-            elif typ == 'dataframe':
-                try:
-                    import pandas as pd
-                except ImportError:
-                    raise SerializationError("Failed to load pandas module")
-                return pd.read_json(val, orient="table")
-            elif typ == 'decimal':
-                return decimal.Decimal(val)
-            elif typ == 'uuid':
-                return uuid.UUID(val)
-            elif typ == 'bytes':
-                return base64.b64decode(val.encode("UTF-8"))
-            return val
-        except Exception as e:
-            raise SerializationError(f"Failed to convert type {typ}: {str(e)}")
+class TempPydanticModel(BaseModel):
+    model_config = {
+        'arbitrary_types_allowed' : True,
+        'extra' : 'allow'
+    }
 
-    @staticmethod
-    def _format_datetime(dt: datetime.datetime) -> str:
-        r = dt.isoformat()
-        if dt.microsecond:
-            r = r[:DATETIME_FORMAT_LENGTH] + r[26:]
-        if r.endswith("+00:00"):
-            r = r[:-6] + "Z"
-        return r
+class MessageSerializer:
+    """Handles message serialization and deserialization."""
 
     @staticmethod
-    def _format_time(t: datetime.time) -> str:
-        r = t.isoformat()
-        if t.microsecond:
-            r = r[:TIME_FORMAT_LENGTH]
-        return r
-
-class IrisJSONEncoder(json.JSONEncoder):
-    """JSONEncoder that handles dates, decimals, UUIDs, etc."""
-    
-    def default(self, obj: Any) -> Any:
+    def _convert_to_json_safe(obj: Any) -> Any:
+        """Convert objects to JSON-safe format."""
         if isinstance(obj, BaseModel):
-            return obj.model_dump()
-        if obj.__class__.__name__ == 'DataFrame':
-            return f'dataframe:{TypeConverter.convert_to_string("dataframe", obj)}'
-        elif isinstance(obj, datetime.datetime):
-            return f'datetime:{TypeConverter.convert_to_string("datetime", obj)}'
-        elif isinstance(obj, datetime.date):
-            return f'date:{TypeConverter.convert_to_string("date", obj)}'
-        elif isinstance(obj, datetime.time):
-            return f'time:{TypeConverter.convert_to_string("time", obj)}'
-        elif isinstance(obj, decimal.Decimal):
-            return f'decimal:{obj}'
-        elif isinstance(obj, uuid.UUID):
-            return f'uuid:{obj}'
-        elif isinstance(obj, bytes):
-            return f'bytes:{TypeConverter.convert_to_string("bytes", obj)}'
-        elif hasattr(obj, '__dict__'):
-            return obj.__dict__
-        return super().default(obj)
-
-class IrisJSONDecoder(json.JSONDecoder):
-    """JSONDecoder that handles special type annotations."""
-    
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super().__init__(object_hook=self.object_hook, *args, **kwargs)
-
-    def object_hook(self, obj: Dict) -> Dict:
-        return {
-            key: self._process_value(value)
-            for key, value in obj.items()
-        }
-                
-    def _process_value(self, value: Any) -> Any:
-        if isinstance(value, str) and TYPE_SEPARATOR in value:
-            typ, val = value.split(TYPE_SEPARATOR, 1)
-            if typ in SUPPORTED_TYPES:
-                return TypeConverter.convert_from_string(typ, val)
-        return value
-
-class MessageSerializer:
-    """Handles message serialization and deserialization."""
+            return obj.model_dump_json()
+        elif is_dataclass(obj):
+            return TempPydanticModel.model_validate(dataclass_to_dict(obj)).model_dump_json()
+        else:
+            raise SerializationError(f"Object {obj} must be a Pydantic model or dataclass")
 
     @staticmethod
     def serialize(message: Any, use_pickle: bool = False) -> iris.cls:
         """Serializes a message to IRIS format."""
-        # Check for PydanticPickleMessage first
-        if isinstance(message, _PydanticPickleMessage):
-            return MessageSerializer._serialize_pickle(message)
-        if isinstance(message, BaseModel):
-            return (MessageSerializer._serialize_pickle(message) 
-                   if use_pickle else MessageSerializer._serialize_json(message))
-        if use_pickle:
+        if isinstance(message, _PydanticPickleMessage) or use_pickle:
             return MessageSerializer._serialize_pickle(message)
         return MessageSerializer._serialize_json(message)
 
-    @staticmethod
-    def deserialize(serial: iris.cls, use_pickle: bool = False) -> Any:
-        """Deserializes a message from IRIS format."""
-        if use_pickle:
-            return MessageSerializer._deserialize_pickle(serial)
-        return MessageSerializer._deserialize_json(serial)
-
-    @staticmethod
-    def _serialize_pickle(message: Any) -> iris.cls:
-        pickle_string = codecs.encode(pickle.dumps(message), "base64").decode()
-        msg = iris.cls('IOP.PickleMessage')._New()
-        msg.classname = f"{message.__class__.__module__}.{message.__class__.__name__}"
-        msg.jstr = _Utils.string_to_stream(pickle_string)
-        return msg
-
     @staticmethod
     def _serialize_json(message: Any) -> iris.cls:
-        if isinstance(message, BaseModel):
-            json_string = json.dumps(message.model_dump(), cls=IrisJSONEncoder, ensure_ascii=False)
-        else:
-            json_string = json.dumps(message, cls=IrisJSONEncoder, ensure_ascii=False)
-            
+        json_string = MessageSerializer._convert_to_json_safe(message)
+        
         msg = iris.cls('IOP.Message')._New()
         msg.classname = f"{message.__class__.__module__}.{message.__class__.__name__}"
         
@@ -178,9 +57,10 @@ def _serialize_json(message: Any) -> iris.cls:
         return msg
 
     @staticmethod
-    def _deserialize_pickle(serial: iris.cls) -> Any:
-        string = _Utils.stream_to_string(serial.jstr)
-        return pickle.loads(codecs.decode(string.encode(), "base64"))
+    def deserialize(serial: iris.cls, use_pickle: bool = False) -> Any:
+        if use_pickle:
+            return MessageSerializer._deserialize_pickle(serial)
+        return MessageSerializer._deserialize_json(serial)
 
     @staticmethod
     def _deserialize_json(serial: iris.cls) -> Any:
@@ -198,14 +78,28 @@ def _deserialize_json(serial: iris.cls) -> Any:
                       if serial.type == 'Stream' else serial.json)
         
         try:
-            json_dict = json.loads(json_string, cls=IrisJSONDecoder)
             if issubclass(msg_class, BaseModel):
-                return msg_class.model_validate(json_dict)
+                return msg_class.model_validate_json(json_string)
+            elif is_dataclass(msg_class):
+                return dataclass_from_dict(msg_class, json.loads(json_string))
             else:
-                return dataclass_from_dict(msg_class, json_dict)
+                raise SerializationError(f"Class {msg_class} must be a Pydantic model or dataclass")
         except Exception as e:
             raise SerializationError(f"Failed to deserialize JSON: {str(e)}")
 
+    @staticmethod
+    def _serialize_pickle(message: Any) -> iris.cls:
+        pickle_string = codecs.encode(pickle.dumps(message), "base64").decode()
+        msg = iris.cls('IOP.PickleMessage')._New()
+        msg.classname = f"{message.__class__.__module__}.{message.__class__.__name__}"
+        msg.jstr = _Utils.string_to_stream(pickle_string)
+        return msg
+
+    @staticmethod
+    def _deserialize_pickle(serial: iris.cls) -> Any:
+        string = _Utils.stream_to_string(serial.jstr)
+        return pickle.loads(codecs.decode(string.encode(), "base64"))
+
     @staticmethod
     def _parse_classname(classname: str) -> tuple[str, str]:
         j = classname.rindex(".")
@@ -214,18 +108,49 @@ def _parse_classname(classname: str) -> tuple[str, str]:
         return classname[:j], classname[j+1:]
 
 def dataclass_from_dict(klass: Type, dikt: Dict) -> Any:
-    """Converts a dictionary to a dataclass instance."""
-    ret = from_dict(klass, dikt, Config(check_types=False))
-    
-    try:
-        fieldtypes = klass.__annotations__
-    except Exception:
-        fieldtypes = {}
-    
-    for key, val in dikt.items():
-        if key not in fieldtypes:
-            setattr(ret, key, val)
-    return ret
+    field_types = {
+        key: val.annotation
+        for key, val in inspect.signature(klass).parameters.items()
+    }
+    processed_dict = {}
+    for key, val in inspect.signature(klass).parameters.items():
+        if key not in dikt and val.default != val.empty:
+            processed_dict[key] = val.default
+            continue
+        
+        value = dikt.get(key)
+        if value is None:
+            processed_dict[key] = None
+            continue
+
+        try:
+            field_type = field_types[key]
+            if field_type != inspect.Parameter.empty:
+                adapter = TypeAdapter(field_type)
+                processed_dict[key] = adapter.validate_python(value)
+            else:
+                processed_dict[key] = value
+        except Exception:
+            processed_dict[key] = value
+
+    instance = klass(
+        **processed_dict
+    )
+    # handle any extra fields
+    for k, v in dikt.items():
+        if k not in processed_dict:
+            setattr(instance, k, v)
+    return instance
+
+def dataclass_to_dict(instance: Any) -> Dict:
+    """Converts a class instance to a dictionary.
+    Handles non attended fields."""
+    dikt = asdict(instance)
+    # assign any extra fields
+    for k, v in vars(instance).items():
+        if k not in dikt:
+            dikt[k] = v
+    return dikt
 
 # Maintain backwards compatibility
 serialize_pickle_message = lambda msg: MessageSerializer.serialize(msg, use_pickle=True)
