[Core] Add type handling registry (Azure#43051)

pvaneck · web-flow · commit 8ba0717ba0dd · 2025-10-09T16:24:11.000-07:00
This registry can be used by serialization and deserialization functions
to enable custom type (de)serialization.

Signed-off-by: Paul Van Eck &lt;paulvaneck@microsoft.com&gt;
diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Added `TypeHandlerRegistry` to `azure.core.serialization` to allow developers to register custom serializers and deserializers for specific types or conditions.  #43051
+
 ### Breaking Changes
 
 ### Bugs Fixed
diff --git a/sdk/core/azure-core/azure/core/serialization.py b/sdk/core/azure-core/azure/core/serialization.py
@@ -5,13 +5,21 @@
 # license information.
 # --------------------------------------------------------------------------
 import base64
+from functools import partial
 from json import JSONEncoder
-from typing import Dict, List, Optional, Union, cast, Any
+from typing import Dict, List, Optional, Union, cast, Any, Type, Callable, Tuple
 from datetime import datetime, date, time, timedelta
 from datetime import timezone
 
 
-__all__ = ["NULL", "AzureJSONEncoder", "is_generated_model", "as_attribute_dict", "attribute_list"]
+__all__ = [
+    "NULL",
+    "AzureJSONEncoder",
+    "is_generated_model",
+    "as_attribute_dict",
+    "attribute_list",
+    "TypeHandlerRegistry",
+]
 TZ_UTC = timezone.utc
 
 
@@ -29,6 +37,164 @@ def __bool__(self) -> bool:
 """
 
 
+class TypeHandlerRegistry:
+    """A registry for custom serializers and deserializers for specific types or conditions."""
+
+    def __init__(self) -> None:
+        self._serializer_types: Dict[Type, Callable] = {}
+        self._deserializer_types: Dict[Type, Callable] = {}
+        self._serializer_predicates: List[Tuple[Callable[[Any], bool], Callable]] = []
+        self._deserializer_predicates: List[Tuple[Callable[[Any], bool], Callable]] = []
+
+        self._serializer_cache: Dict[Type, Optional[Callable]] = {}
+        self._deserializer_cache: Dict[Type, Optional[Callable]] = {}
+
+    def register_serializer(
+        self, condition: Union[Type, Callable[[Any], bool]]
+    ) -> Callable[[Callable[[Any], Dict[str, Any]]], Callable[[Any], Dict[str, Any]]]:
+        """Decorator to register a serializer.
+
+        The handler function is expected to take a single argument, the object to serialize,
+        and return a dictionary representation of that object.
+
+        Examples:
+
+        .. code-block:: python
+
+            @registry.register_serializer(CustomModel)
+            def serialize_single_type(value: CustomModel) -> dict:
+                return value.to_dict()
+
+            @registry.register_serializer(lambda x: isinstance(x, BaseModel))
+            def serialize_with_condition(value: BaseModel) -> dict:
+                return value.to_dict()
+
+            # Called manually for a specific type
+            def custom_serializer(value: CustomModel) -> Dict[str, Any]:
+                return {"custom": value.custom}
+
+            registry.register_serializer(CustomModel)(custom_serializer)
+
+        :param condition: A type or a callable predicate function that takes an object and returns a bool.
+        :type condition: Union[Type, Callable[[Any], bool]]
+        :return: A decorator that registers the handler function.
+        :rtype: Callable[[Callable[[Any], Dict[str, Any]]], Callable[[Any], Dict[str, Any]]]
+        :raises TypeError: If the condition is neither a type nor a callable.
+        """
+
+        def decorator(handler_func: Callable[[Any], Dict[str, Any]]) -> Callable[[Any], Dict[str, Any]]:
+            if isinstance(condition, type):
+                self._serializer_types[condition] = handler_func
+            elif callable(condition):
+                self._serializer_predicates.append((condition, handler_func))
+            else:
+                raise TypeError("Condition must be a type or a callable predicate function.")
+
+            self._serializer_cache.clear()
+            return handler_func
+
+        return decorator
+
+    def register_deserializer(
+        self, condition: Union[Type, Callable[[Any], bool]]
+    ) -> Callable[[Callable[[Type, Dict[str, Any]], Any]], Callable[[Type, Dict[str, Any]], Any]]:
+        """Decorator to register a deserializer.
+
+        The handler function is expected to take two arguments: the target type and the data dictionary,
+        and return an instance of the target type.
+
+        Examples:
+
+        .. code-block:: python
+
+            @registry.register_deserializer(CustomModel)
+            def deserialize_single_type(cls: Type[CustomModel], data: dict) -> CustomModel:
+                return cls(**data)
+
+            @registry.register_deserializer(lambda t: issubclass(t, BaseModel))
+            def deserialize_with_condition(cls: Type[BaseModel], data: dict) -> BaseModel:
+                return cls(**data)
+
+            # Called manually for a specific type
+            def custom_deserializer(cls: Type[CustomModel], data: Dict[str, Any]) -> CustomModel:
+                return cls(custom=data["custom"])
+
+            registry.register_deserializer(CustomModel)(custom_deserializer)
+
+        :param condition: A type or a callable predicate function that takes an object and returns a bool.
+        :type condition: Union[Type, Callable[[Any], bool]]
+        :return: A decorator that registers the handler function.
+        :rtype: Callable[[Callable[[Type, Dict[str, Any]], Any]], Callable[[Type, Dict[str, Any]], Any]]
+        :raises TypeError: If the condition is neither a type nor a callable.
+        """
+
+        def decorator(handler_func: Callable[[Type, Dict[str, Any]], Any]) -> Callable[[Type, Dict[str, Any]], Any]:
+            if isinstance(condition, type):
+                self._deserializer_types[condition] = handler_func
+            elif callable(condition):
+                self._deserializer_predicates.append((condition, handler_func))
+            else:
+                raise TypeError("Condition must be a type or a callable predicate function.")
+
+            self._deserializer_cache.clear()
+            return handler_func
+
+        return decorator
+
+    def get_serializer(self, obj: Any) -> Optional[Callable]:
+        """Gets the appropriate serializer for an object.
+
+        It first checks the type dictionary for a direct type match.
+        If no match is found, it iterates through the predicate list to find a match.
+
+        Results of the lookup are cached for performance based on the object's type.
+
+        :param obj: The object to serialize.
+        :type obj: any
+        :return: The serializer function if found, otherwise None.
+        :rtype: Optional[Callable]
+        """
+        obj_type = type(obj)
+        if obj_type in self._serializer_cache:
+            return self._serializer_cache[obj_type]
+
+        handler = self._serializer_types.get(type(obj))
+        if not handler:
+            for predicate, pred_handler in self._serializer_predicates:
+                if predicate(obj):
+                    handler = pred_handler
+                    break
+
+        self._serializer_cache[obj_type] = handler
+        return handler
+
+    def get_deserializer(self, cls: Type) -> Optional[Callable]:
+        """Gets the appropriate deserializer for a class.
+
+        It first checks the type dictionary for a direct type match.
+        If no match is found, it iterates through the predicate list to find a match.
+
+        Results of the lookup are cached for performance based on the class.
+
+        :param cls: The class to deserialize.
+        :type cls: type
+        :return: The deserializer function wrapped with the class if found, otherwise None.
+        :rtype: Optional[Callable]
+        """
+        if cls in self._deserializer_cache:
+            return self._deserializer_cache[cls]
+
+        handler = self._deserializer_types.get(cls)
+        if not handler:
+            for predicate, pred_handler in self._deserializer_predicates:
+                if predicate(cls):
+                    handler = pred_handler
+                    break
+
+        self._deserializer_cache[cls] = partial(handler, cls) if handler else None
+        return self._deserializer_cache[cls]
+
+
 def _timedelta_as_isostr(td: timedelta) -> str:
     """Converts a datetime.timedelta object into an ISO 8601 formatted string, e.g. 'P4DT12H30M05S'
 
diff --git a/sdk/core/azure-core/tests/specs_sdk/modeltypes/modeltypes/_utils/model_base.py b/sdk/core/azure-core/tests/specs_sdk/modeltypes/modeltypes/_utils/model_base.py
@@ -17,6 +17,7 @@
 import base64
 import re
 import typing
+import traceback
 import enum
 import email.utils
 from datetime import datetime, date, time, timedelta, timezone
@@ -28,7 +29,7 @@
 from azure.core.exceptions import DeserializationError
 from azure.core import CaseInsensitiveEnumMeta
 from azure.core.pipeline import PipelineResponse
-from azure.core.serialization import _Null
+from azure.core.serialization import _Null, TypeHandlerRegistry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -38,6 +39,9 @@
 _T = typing.TypeVar("_T")
 
 
+TYPE_HANDLER_REGISTRY = TypeHandlerRegistry()
+
+
 def _timedelta_as_isostr(td: timedelta) -> str:
     """Converts a datetime.timedelta object into an ISO 8601 formatted string, e.g. 'P4DT12H30M05S'
 
@@ -161,6 +165,10 @@ def default(self, o):  # pylint: disable=too-many-return-statements
             except AttributeError:
                 # This will be raised when it hits value.total_seconds in the method above
                 pass
+
+            custom_serializer = TYPE_HANDLER_REGISTRY.get_serializer(o)
+            if custom_serializer:
+                return custom_serializer(o)
             return super(SdkJSONEncoder, self).default(o)
 
 
@@ -481,6 +489,7 @@ def _is_model(obj: typing.Any) -> bool:
 
 
 def _serialize(o, format: typing.Optional[str] = None):  # pylint: disable=too-many-return-statements
+
     if isinstance(o, list):
         return [_serialize(x, format) for x in o]
     if isinstance(o, dict):
@@ -510,6 +519,12 @@ def _serialize(o, format: typing.Optional[str] = None):  # pylint: disable=too-m
     except AttributeError:
         # This will be raised when it hits value.total_seconds in the method above
         pass
+
+    # Check if there's a custom serializer for the type
+    custom_serializer = TYPE_HANDLER_REGISTRY.get_serializer(o)
+    if custom_serializer:
+        return custom_serializer(o)
+
     return o
 
 
@@ -886,6 +901,10 @@ def _deserialize_default(
     if get_deserializer(annotation, rf):
         return functools.partial(_deserialize_default, get_deserializer(annotation, rf))
 
+    deserializer = TYPE_HANDLER_REGISTRY.get_deserializer(annotation)
+    if deserializer:
+        return deserializer
+
     return functools.partial(_deserialize_default, annotation)
 
 
diff --git a/sdk/core/azure-core/tests/test_serialization.py b/sdk/core/azure-core/tests/test_serialization.py
