add support to dynamic model deducation

Author: aradyaron

aws_lambda_powertools/utilities/idempotency/exceptions.py

@@ -71,3 +71,15 @@ class IdempotencyKeyError(BaseError):
     """
     Payload does not contain an idempotent key
     """
+
+
+class IdempotencyModelTypeError(BaseError):
+    """
+    Model type does not match expected type
+    """
+
+
+class IdempotencyNoSerializationModelError(BaseError):
+    """
+    No model was supplied to the serializer
+    """

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -4,7 +4,8 @@
 import functools
 import logging
 import os
-from typing import Any, Callable, Dict, Optional, cast
+from inspect import isclass
+from typing import Any, Callable, Dict, Optional, Type, Union, cast
 
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
 from aws_lambda_powertools.shared import constants
@@ -14,7 +15,10 @@
 from aws_lambda_powertools.utilities.idempotency.persistence.base import (
     BasePersistenceLayer,
 )
-from aws_lambda_powertools.utilities.idempotency.serialization.base import BaseIdempotencySerializer
+from aws_lambda_powertools.utilities.idempotency.serialization.base import (
+    BaseIdempotencyModelSerializer,
+    BaseIdempotencySerializer,
+)
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 logger = logging.getLogger(__name__)
@@ -86,7 +90,7 @@ def idempotent_function(
     data_keyword_argument: str,
     persistence_store: BasePersistenceLayer,
     config: Optional[IdempotencyConfig] = None,
-    output_serializer: Optional[BaseIdempotencySerializer] = None,
+    output_serializer: Optional[Union[BaseIdempotencySerializer, Type[BaseIdempotencyModelSerializer]]] = None,
 ) -> Any:
     """
     Decorator to handle idempotency of any function
@@ -101,9 +105,11 @@ def idempotent_function(
         Instance of BasePersistenceLayer to store data
     config: IdempotencyConfig
         Configuration
-    output_serializer: Optional[BaseIdempotencySerializer]
+    output_serializer: Optional[Union[BaseIdempotencySerializer, Type[BaseIdempotencyModelSerializer]]]
             Serializer to transform the data to and from a dictionary.
-            If not supplied, no serialization is done via the NoOpSerializer
+            If not supplied, no serialization is done via the NoOpSerializer.
+            In case a serializer of type inheriting BaseIdempotencyModelSerializer] is given,
+            the serializer is deduced from the function return type.
     Examples
     --------
     **Processes an order in an idempotent manner**
@@ -131,6 +137,9 @@ def process_order(customer_id: str, order: dict, **kwargs):
                 output_serializer=output_serializer,
             ),
         )
+    if isclass(output_serializer) and issubclass(output_serializer, BaseIdempotencyModelSerializer):
+        # instantiate an instance of the serializer class
+        output_serializer = output_serializer.instantiate(function.__annotations__.get("return", None))
 
     config = config or IdempotencyConfig()
 

aws_lambda_powertools/utilities/idempotency/serialization/base.py

@@ -17,3 +17,31 @@ def to_dict(self, data: Any) -> Dict:
     @abstractmethod
     def from_dict(self, data: Dict) -> Any:
         pass
+
+
+class BaseIdempotencyModelSerializer(BaseIdempotencySerializer):
+    """
+    Abstract Base Class for Idempotency serialization layer, for using a model as data object representation.
+    """
+
+    @classmethod
+    @abstractmethod
+    def instantiate(cls, model_type: Any) -> BaseIdempotencySerializer:
+        """
+        Instantiate the serializer from the given model type.
+        In case there is the model_type is unknown, None will be sent to the method.
+        It's on the implementer to verify that:
+        - None is handled
+        - A model type not matching the expected types is handled
+
+        Parameters
+        ----------
+        model_type: Any
+            The model type to instantiate the class for
+
+        Returns
+        -------
+        BaseIdempotencySerializer
+            Instance of the serializer class
+        """
+        pass

aws_lambda_powertools/utilities/idempotency/serialization/pydantic.py

@@ -1,13 +1,20 @@
-from typing import Dict, Type
+from typing import Any, Dict, Type
 
 from pydantic import BaseModel
 
-from aws_lambda_powertools.utilities.idempotency.serialization.base import BaseIdempotencySerializer
+from aws_lambda_powertools.utilities.idempotency.exceptions import (
+    IdempotencyModelTypeError,
+    IdempotencyNoSerializationModelError,
+)
+from aws_lambda_powertools.utilities.idempotency.serialization.base import (
+    BaseIdempotencyModelSerializer,
+    BaseIdempotencySerializer,
+)
 
 Model = BaseModel
 
 
-class PydanticSerializer(BaseIdempotencySerializer):
+class PydanticSerializer(BaseIdempotencyModelSerializer):
     def __init__(self, model: Type[Model]):
         """
         Parameters
@@ -28,3 +35,13 @@ def from_dict(self, data: Dict) -> Model:
             # Support for pydantic V2
             return self.__model.model_validate(data)  # type: ignore[unused-ignore,attr-defined]
         return self.__model.parse_obj(data)
+
+    @classmethod
+    def instantiate(cls, model_type: Any) -> BaseIdempotencySerializer:
+        if model_type is None:
+            raise IdempotencyNoSerializationModelError("No serialization model was supplied")
+
+        if not issubclass(model_type, BaseModel):
+            raise IdempotencyModelTypeError("Model type is not inherited from pydantic BaseModel")
+
+        return cls(model=model_type)

tests/functional/idempotency/test_idempotency.py

@@ -1278,7 +1278,8 @@ def record_handler(record):
     assert from_dict_called is False, "in case response is None, from_dict should not be called"
 
 
-def test_idempotent_function_serialization_pydantic():
+@pytest.mark.parametrize("output_serializer_type", ["explicit", "deduced"])
+def test_idempotent_function_serialization_pydantic(output_serializer_type: str):
     # GIVEN
     config = IdempotencyConfig(use_local_cache=True)
     mock_event = {"customer_id": "fake", "transaction_id": "fake-id"}
@@ -1293,13 +1294,18 @@ class PaymentOutput(BaseModel):
         customer_id: str
         transaction_id: str
 
+    if output_serializer_type == "explicit":
+        output_serializer = PydanticSerializer(
+            model=PaymentOutput,
+        )
+    else:
+        output_serializer = PydanticSerializer
+
     @idempotent_function(
         data_keyword_argument="payment",
         persistence_store=persistence_layer,
         config=config,
-        output_serializer=PydanticSerializer(
-            model=PaymentOutput,
-        ),
+        output_serializer=output_serializer,
     )
     def collect_payment(payment: PaymentInput) -> PaymentOutput:
         return PaymentOutput(**payment.dict())