diff --git a/amqp/adapter/pydantic_serializer.py b/amqp/adapter/pydantic_serializer.py index 1fdbf72..27a61ae 100644 --- a/amqp/adapter/pydantic_serializer.py +++ b/amqp/adapter/pydantic_serializer.py @@ -2,10 +2,12 @@ This module provides serialization and deserialization functions for Pydantic models that extend BaseModel """ +import importlib import json import uuid from datetime import datetime -from typing import Any, Dict +from types import UnionType +from typing import Any, Dict, Union, get_args, get_origin from urllib.parse import parse_qs, urlparse from pydantic import BaseModel @@ -26,8 +28,12 @@ def serialize_object(obj: BaseModel) -> str: A string representation of the object. """ class_name = obj.__class__.__name__ + module_name = obj.__class__.__module__ obj_dict = {} - for key, value in obj.model_dump().items(): + for key, _ in obj.model_dump().items(): + # here we use the original value + # the one returned by model_dump will convert everything into dict recursively + value = getattr(obj, key) if isinstance(value, uuid.UUID): obj_dict[key] = str(value) elif isinstance(value, datetime): @@ -47,8 +53,21 @@ def serialize_object(obj: BaseModel) -> str: else: obj_dict[key] = value - dict_string = ",".join(f"{k}:{v}" for k, v in obj_dict.items()) - return f"{class_name}:{{{dict_string}}}" + dict_string = json.dumps(obj_dict) + return f"{module_name}.{class_name}:{dict_string}" + + +def try_deserialize_object(data: str) -> BaseModel | str: + """ + Try deserialize object, if failed, return the input string as-is. + """ + if not isinstance(data, str): + return data + try: + return deserialize_object(data) + except Exception as e: + logging_debug(f"Failed to deserialize object: {e}") + return data def deserialize_object(data: str) -> BaseModel: @@ -63,93 +82,54 @@ def deserialize_object(data: str) -> BaseModel: Returns: An R2RSerializable object. """ - class_name, dict_str = data.split(":", 1) - dict_str = dict_str.strip("{}") - pairs = dict_str.split(",") if dict_str else [] + class_path, dict_str = data.split(":", 1) + module_path, class_name = class_path.rsplit(".", 1) - obj_dict: Dict[str, Any] = {} - for pair in pairs: - if not pair or ":" not in pair: - continue - k, v = pair.split(":", 1) - v = v.strip() # important - obj_dict[k] = _convert_value(v) + obj_dict: Dict[str, Any] = json.loads(dict_str) - logging_debug(f"Deserialized object: {class_name}") + logging_debug(f"Deserialized object: {class_path}") if "[" in class_name: class_name = class_name.split("[")[0] logging_debug(f"trying object: {class_name}") - cls = globals().get(class_name) + module = importlib.import_module(module_path) + cls = getattr(module, class_name) if cls is None: raise ValueError(f"Unknown class: {class_name}") + # process nested model, dict and list + for key, value in obj_dict.items(): + key_type: type[Any] = cls.model_fields[key].annotation + # handle union with none, this will give the true type + # for example, return `A` if the type is `A | None` + if get_origin(key_type) is UnionType or get_origin(key_type) is Union: + union_types = get_args(key_type) + if len(union_types) > 2: + # if there are multiple unions, like `A | B | None` + # then we have no way to figure out which one to use + raise ValueError(f"Unsupported Union type: {key_type}") + # return the none-None one as the type for processing + if union_types[0] is type(None): + key_type = union_types[1] + else: + key_type = union_types[0] + # skip if the value is None + if value is None: + continue + + if key_type == uuid.UUID: + obj_dict[key] = uuid.UUID(value) + elif key_type == datetime: + obj_dict[key] = datetime.fromisoformat(value) + elif key_type == dict: + obj_dict[key] = value + elif issubclass(key_type, BaseModel): + obj_dict[key] = deserialize_object(value) + elif key_type == list: + obj_dict[key] = list(map(try_deserialize_object, value)) + return cls(**obj_dict) -def _convert_value(v: str) -> Any: - """ - Helper function to convert a string to its appropriate Python type. - """ - if v == "None": - return None - elif _is_uuid_string(v): - return uuid.UUID(v) - elif _is_datetime_string(v): - return datetime.fromisoformat(v) - elif _is_int_string(v): - return int(v) - elif ( - v.startswith("{") - and v.endswith("}") - and ":" in v - and "{" not in v[1:-1] - and "}" not in v[1:-1] - ): # nested object - # check if it is a json dict - try: - json.loads(v) - return json.loads(v) - except json.JSONDecodeError: - return deserialize_object(v) # recursive call - elif v.startswith("[") and v.endswith("]"): # list - v = v[1:-1] # remove brackets - items = v.split(",") if v else [] - deserialized_list = [] - for item in items: - item = item.strip() - deserialized_list.append(_convert_value(item)) # recursive call - return deserialized_list - else: - return v # Keep as string - - -def _is_uuid_string(s: str) -> bool: - """Helper to check if a string is a valid UUID.""" - try: - uuid.UUID(s) - return True - except ValueError: - return False - - -def _is_datetime_string(s: str) -> bool: - """Helper to check if a string is a valid datetime in ISO format""" - try: - datetime.fromisoformat(s) - return True - except ValueError: - return False - - -def _is_int_string(s: str) -> bool: - """Helper to check if a string is a valid int.""" - try: - int(s) - return True - except ValueError: - return False - - def parse_url_query(url): """ Parses a HTTP GET URL and returns (context_path, query_params_dict) diff --git a/pyproject.toml b/pyproject.toml index c86f5b5..60fc23a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,5 +49,6 @@ dev = [ "black", "isort", "flake8", - "pre-commit" + "pre-commit", + "pylint" ] diff --git a/tests/adapter/test_pydantic_serializer.py b/tests/adapter/test_pydantic_serializer.py index 5e2bbfd..af2b9ff 100644 --- a/tests/adapter/test_pydantic_serializer.py +++ b/tests/adapter/test_pydantic_serializer.py @@ -1,38 +1,132 @@ +from datetime import datetime from unittest import TestCase +from uuid import UUID -from amqp.adapter.pydantic_serializer import parse_url_query +from pydantic import BaseModel + +from amqp.adapter.pydantic_serializer import ( + deserialize_object, + parse_url_query, + serialize_object, +) + + +class NestedModel(BaseModel): + """ + Nested model for testing the pydantic serializer. + """ + + value: str + number: int + + +class TestModel(BaseModel): + """ + Test model for testing the pydantic serializer. + """ + + name: str + age: int + uuid_field: UUID + datetime_field: datetime + nested_model: NestedModel | None = None + list_field: list = [] class PydanticSerializerTest(TestCase): + + def setUp(self): + """ + Prepare the test data. + """ + self.test_uuid = UUID("12345678-1234-5678-1234-567812345678") + self.test_datetime = datetime(2024, 1, 1, 12, 0, 0) + self.nested_model = NestedModel(value="test", number=42) + self.test_model = TestModel( + name="test", + age=25, + uuid_field=self.test_uuid, + datetime_field=self.test_datetime, + nested_model=self.nested_model, + list_field=[1, 2, 3], + ) + def test_serialize_object(self): - assert True + serialized = serialize_object(self.test_model) + self.assertIsInstance(serialized, str) + self.assertIn("TestModel", serialized) + self.assertIn("test", serialized) + self.assertIn("25", serialized) + self.assertIn(str(self.test_uuid), serialized) + self.assertIn(self.test_datetime.isoformat(), serialized) def test_deserialize_object(self): - assert True + serialized = serialize_object(self.test_model) + deserialized = deserialize_object(serialized) - def test__convert_value(self): - assert True + self.assertIsInstance(deserialized, TestModel) + self.assertEqual(deserialized.name, "test") + self.assertEqual(deserialized.age, 25) + self.assertEqual(deserialized.uuid_field, self.test_uuid) + self.assertEqual(deserialized.datetime_field, self.test_datetime) + self.assertIsInstance(deserialized.nested_model, NestedModel) + self.assertEqual(deserialized.nested_model.value, "test") + self.assertEqual(deserialized.nested_model.number, 42) + self.assertEqual(deserialized.list_field, [1, 2, 3]) - def test__is_uuid_string(self): - assert True + def test_serialize_deserialize_with_none(self): + model = TestModel( + name="test", + age=25, + uuid_field=self.test_uuid, + datetime_field=self.test_datetime, + nested_model=None, + list_field=[], + ) + serialized = serialize_object(model) + deserialized = deserialize_object(serialized) + self.assertIsNone(deserialized.nested_model) - def test__is_datetime_string(self): - assert True + def test_serialize_deserialize_with_list_of_models(self): + nested_models = [NestedModel(value=f"test{i}", number=i) for i in range(3)] + model = TestModel( + name="test", + age=25, + uuid_field=self.test_uuid, + datetime_field=self.test_datetime, + nested_model=None, + list_field=nested_models, + ) + serialized = serialize_object(model) + deserialized = deserialize_object(serialized) + self.assertEqual(len(deserialized.list_field), 3) + for i, item in enumerate(deserialized.list_field): + self.assertIsInstance(item, NestedModel) + self.assertEqual(item.value, f"test{i}") + self.assertEqual(item.number, i) - def test__is_int_string(self): - assert True + def test_parse_url_query_no_params(self): + test_url = "http://example.com/path" + path, params = parse_url_query(test_url) - def test_serialize_to_json(self): - assert True + self.assertEqual(path, "/path") + self.assertEqual(params, {}) + + def test_parse_url_query_single_value(self): + test_url = "http://example.com/path?foo=1" + path, params = parse_url_query(test_url) + + self.assertEqual(path, "/path") + self.assertEqual(params["foo"], "1") def test_parse_url_query(self): path, args = parse_url_query( "http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741?response_type=JsonFile" ) - assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741" - assert args == {"response_type": "JsonFile"} + self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741") + self.assertEqual(args, {"response_type": "JsonFile"}) path, args = parse_url_query( "http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741" ) - assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741" - assert args == {} + self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741") + self.assertEqual(args, {})