diff --git a/aws_lambda_powertools/utilities/idempotency/idempotency.py b/aws_lambda_powertools/utilities/idempotency/idempotency.py index 62a8994889c..dc05a30a239 100644 --- a/aws_lambda_powertools/utilities/idempotency/idempotency.py +++ b/aws_lambda_powertools/utilities/idempotency/idempotency.py @@ -7,9 +7,10 @@ import functools import logging import os +import sys import warnings -from inspect import isclass -from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast +from inspect import isclass, unwrap +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast, get_type_hints from aws_lambda_powertools.middleware_factory import lambda_handler_decorator from aws_lambda_powertools.shared import constants @@ -145,6 +146,34 @@ def handler(event, context): idempotent = cast(_IdempotentDecorator, _idempotent) +def _get_return_annotation(function: Callable) -> Any: + """Resolve only the return annotation, preserving evaluated types and custom metadata.""" + if sys.version_info >= (3, 14): + from annotationlib import Format, get_annotations + + # Deferred annotations can contain parameter names that only exist during type checking. + annotation = get_annotations(function, format=Format.FORWARDREF).get("return") + else: + annotation = getattr(function, "__annotations__", {}).get("return") + + if not isinstance(annotation, str): + return annotation + + def return_type(): ... + + # A separate function avoids evaluating input annotations or copying __no_type_check__. + return_type.__annotations__ = {"return": annotation} + try: + return get_type_hints( + return_type, + globalns=getattr(unwrap(function), "__globals__", {}), + include_extras=True, + )["return"] + except Exception: + # Let the serializer report an unsupported model, as for other invalid annotations. + return annotation + + def idempotent_function( function: AnyCallableT | None = None, *, @@ -208,7 +237,7 @@ def process_order(customer_id: str, order: dict, **kwargs): 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)) + output_serializer = output_serializer.instantiate(_get_return_annotation(function)) config = config or IdempotencyConfig() diff --git a/aws_lambda_powertools/utilities/idempotency/serialization/pydantic.py b/aws_lambda_powertools/utilities/idempotency/serialization/pydantic.py index 59d67f5f3c9..0653d1b5127 100644 --- a/aws_lambda_powertools/utilities/idempotency/serialization/pydantic.py +++ b/aws_lambda_powertools/utilities/idempotency/serialization/pydantic.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import Any +from inspect import isclass +from typing import Any, get_origin from pydantic import BaseModel @@ -41,7 +42,7 @@ 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): + if get_origin(model_type) is not None or not isclass(model_type) or not issubclass(model_type, BaseModel): raise IdempotencyModelTypeError("Model type is not inherited from pydantic BaseModel") return cls(model=model_type) diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md index da7e70ad89c..247b15f7eb0 100644 --- a/docs/utilities/idempotency.md +++ b/docs/utilities/idempotency.md @@ -223,13 +223,17 @@ The output serializer supports any JSON serializable data, **Python Dataclasses* Function returns must be annotated with a single type, optionally wrapped in `Optional` or `Union` with `None`. + Inference supports `from __future__ import annotations`. Define the output model at module scope so its name can be resolved when the function is decorated. + If the model is local to another function, pass an explicit serializer instance, such as `PydanticSerializer(model=OrderOutput)` or `DataclassSerializer(model=OrderOutput)`. + Containers such as `list[OrderOutput]` are not supported by these model serializers and raise `IdempotencyModelTypeError`. + === "Pydantic" Use `PydanticSerializer` to automatically serialize what's retrieved from the persistent storage based on the return type annotated. === "Inferring via the return type" - ```python hl_lines="8 27 35 38 48" + ```python hl_lines="1 10 29 37 40 50" --8<-- "examples/idempotency/src/working_with_pydantic_deduced_output_serializer.py" ``` @@ -255,7 +259,7 @@ The output serializer supports any JSON serializable data, **Python Dataclasses* === "Inferring via the return type" - ```python hl_lines="9 30 38 41 51" + ```python hl_lines="1 11 32 40 43 53" --8<-- "examples/idempotency/src/working_with_dataclass_deduced_output_serializer.py" ``` diff --git a/examples/idempotency/src/working_with_dataclass_deduced_output_serializer.py b/examples/idempotency/src/working_with_dataclass_deduced_output_serializer.py index e6f74cb8f9a..b262dc4e84a 100644 --- a/examples/idempotency/src/working_with_dataclass_deduced_output_serializer.py +++ b/examples/idempotency/src/working_with_dataclass_deduced_output_serializer.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import os from dataclasses import dataclass diff --git a/examples/idempotency/src/working_with_pydantic_deduced_output_serializer.py b/examples/idempotency/src/working_with_pydantic_deduced_output_serializer.py index b904a5ad670..f4137bbfa6c 100644 --- a/examples/idempotency/src/working_with_pydantic_deduced_output_serializer.py +++ b/examples/idempotency/src/working_with_pydantic_deduced_output_serializer.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import os from aws_lambda_powertools.utilities.idempotency import ( diff --git a/tests/functional/idempotency/_pydantic/test_idempotency_return_annotation_compatibility.py b/tests/functional/idempotency/_pydantic/test_idempotency_return_annotation_compatibility.py new file mode 100644 index 00000000000..063d8cebdd5 --- /dev/null +++ b/tests/functional/idempotency/_pydantic/test_idempotency_return_annotation_compatibility.py @@ -0,0 +1,140 @@ +import functools +import sys +from typing import TYPE_CHECKING, Annotated, get_args, get_origin, no_type_check + +import pytest + +from aws_lambda_powertools.utilities.idempotency import IdempotencyConfig, idempotent_function +from aws_lambda_powertools.utilities.idempotency.serialization.base import BaseIdempotencyModelSerializer +from aws_lambda_powertools.utilities.idempotency.serialization.dataclass import DataclassSerializer +from aws_lambda_powertools.utilities.idempotency.serialization.pydantic import PydanticSerializer +from tests.functional.idempotency._pydantic.test_idempotency_with_postponed_annotations import ( + MOCK_EVENT, + NoOpPersistenceLayer, + PaymentOutput, + PaymentOutputDataclass, +) + +if TYPE_CHECKING: + from collections.abc import Mapping as InputOnlyAvailableToTypeChecker + + +class MetadataSerializer(BaseIdempotencyModelSerializer): + @classmethod + def instantiate(cls, model_type): + assert get_origin(model_type) is Annotated + model, metadata = get_args(model_type) + assert metadata == "v1" + return PydanticSerializer(model=model) + + def to_dict(self, data): ... + + def from_dict(self, data): ... + + +def replay(function, serializer, expected): + config = IdempotencyConfig(use_local_cache=True) + + # This fake Lambda context avoids depending on DynamoDB or warning about missing context. + class Context: + def get_remaining_time_in_millis(self): + return 1000 + + config.register_lambda_context(Context()) + decorated = idempotent_function( + function, + data_keyword_argument="payment", + persistence_store=NoOpPersistenceLayer(), + config=config, + output_serializer=serializer, + ) + assert decorated(payment=MOCK_EVENT) == expected + assert decorated(payment=MOCK_EVENT) == expected + + +@pytest.mark.parametrize("postponed", [False, True]) +def test_custom_serializer_keeps_annotated_metadata(postponed): + calls = [] + + def collect_payment(payment: dict) -> Annotated[PaymentOutput, "v1"]: + calls.append(payment) + return PaymentOutput(**payment) + + if postponed: + collect_payment.__annotations__["return"] = "Annotated[PaymentOutput, 'v1']" + + replay(collect_payment, MetadataSerializer, PaymentOutput(**MOCK_EVENT)) + assert len(calls) == 1 + + +@pytest.mark.parametrize("postponed", [False, True]) +def test_no_type_check_keeps_return_model(postponed): + calls = [] + + @no_type_check + def collect_payment(payment: dict) -> PaymentOutput: + calls.append(payment) + return PaymentOutput(**payment) + + if postponed: + collect_payment.__annotations__["return"] = "PaymentOutput" + + replay(collect_payment, PydanticSerializer, PaymentOutput(**MOCK_EVENT)) + assert len(calls) == 1 + + +@pytest.mark.parametrize( + "model,serializer", + [(PaymentOutput, PydanticSerializer), (PaymentOutputDataclass, DataclassSerializer)], +) +def test_unresolved_input_annotation_does_not_block_return_model(model, serializer): + calls = [] + + def collect_payment(payment: "InputOnlyAvailableToTypeChecker"): + calls.append(payment) + return model(**payment) + + collect_payment.__annotations__["return"] = model.__name__ + replay(collect_payment, serializer, model(**MOCK_EVENT)) + assert len(calls) == 1 + + +def test_wrapped_function_resolves_return_in_original_namespace(): + calls = [] + + def collect_payment(payment: dict) -> "PaymentOutput": + calls.append(payment) + return PaymentOutput(**payment) + + # A wrapper from a different module has no PaymentOutput in its globals. + namespace = {"functools": functools, "original": collect_payment} + exec( + compile( + "@functools.wraps(original)\ndef wrapper(*args, **kwargs):\n return original(*args, **kwargs)\n", + "", + "exec", + dont_inherit=True, + ), + namespace, + ) + replay(namespace["wrapper"], PydanticSerializer, PaymentOutput(**MOCK_EVENT)) + assert len(calls) == 1 + + +@pytest.mark.skipif(sys.version_info < (3, 14), reason="Native deferred annotations require Python 3.14") +def test_native_deferred_annotations_allow_unresolved_input(): + calls = [] + namespace = {"PaymentOutput": PaymentOutput, "calls": calls} + exec( + compile( + "def collect_payment(payment: InputOnlyAvailableToTypeChecker) -> PaymentOutput:\n" + " calls.append(payment)\n" + " return PaymentOutput(**payment)\n", + "", + "exec", + dont_inherit=True, + ), + namespace, + ) + replay(namespace["collect_payment"], PydanticSerializer, PaymentOutput(**MOCK_EVENT)) + assert len(calls) == 1 diff --git a/tests/functional/idempotency/_pydantic/test_idempotency_with_postponed_annotations.py b/tests/functional/idempotency/_pydantic/test_idempotency_with_postponed_annotations.py new file mode 100644 index 00000000000..e5d9e092003 --- /dev/null +++ b/tests/functional/idempotency/_pydantic/test_idempotency_with_postponed_annotations.py @@ -0,0 +1,155 @@ +# Postponed annotations turn every annotation in this module into a string, as in customer code that uses +# `from __future__ import annotations`. Models are defined at module level, where customers define them. +from __future__ import annotations + +import dataclasses + +import pytest +from pydantic import BaseModel + +from aws_lambda_powertools.utilities.idempotency import ( + IdempotencyConfig, + idempotent_function, +) +from aws_lambda_powertools.utilities.idempotency.exceptions import ( + IdempotencyModelTypeError, +) +from aws_lambda_powertools.utilities.idempotency.persistence.base import ( + BasePersistenceLayer, + DataRecord, +) +from aws_lambda_powertools.utilities.idempotency.serialization.dataclass import ( + DataclassSerializer, +) +from aws_lambda_powertools.utilities.idempotency.serialization.pydantic import ( + PydanticSerializer, +) + + +class PaymentInput(BaseModel): + customer_id: str + transaction_id: str + + +class PaymentOutput(BaseModel): + customer_id: str + transaction_id: str + + +@dataclasses.dataclass +class PaymentOutputDataclass: + customer_id: str + transaction_id: str + + +class NoOpPersistenceLayer(BasePersistenceLayer): + # Responses are replayed from the local cache, so the store itself doesn't need to keep anything + def _put_record(self, data_record: DataRecord) -> None: ... + + def _update_record(self, data_record: DataRecord) -> None: ... + + def _get_record(self, idempotency_key) -> DataRecord: ... + + def _delete_record(self, data_record: DataRecord) -> None: ... + + +MOCK_EVENT = {"customer_id": "fake", "transaction_id": "fake-id"} + + +def test_idempotent_function_pydantic_serializer_deduced_from_postponed_annotation(): + # GIVEN a function annotated with a Pydantic model under postponed annotations + @idempotent_function( + data_keyword_argument="payment", + persistence_store=NoOpPersistenceLayer(), + config=IdempotencyConfig(use_local_cache=True), + output_serializer=PydanticSerializer, + ) + def collect_payment(payment: PaymentInput) -> PaymentOutput: + return PaymentOutput(**payment.model_dump()) + + # WHEN calling it twice with the same payload + payment = PaymentInput(**MOCK_EVENT) + first_call = collect_payment(payment=payment) + second_call = collect_payment(payment=payment) + + # THEN the replayed response is rebuilt as the annotated model + assert isinstance(first_call, PaymentOutput) + assert isinstance(second_call, PaymentOutput) + assert second_call == first_call + + +def test_idempotent_function_pydantic_serializer_deduced_from_postponed_optional_annotation(): + # GIVEN a function annotated with an optional Pydantic model under postponed annotations + @idempotent_function( + data_keyword_argument="payment", + persistence_store=NoOpPersistenceLayer(), + config=IdempotencyConfig(use_local_cache=True), + output_serializer=PydanticSerializer, + ) + def collect_payment(payment: PaymentInput) -> PaymentOutput | None: + return PaymentOutput(**payment.model_dump()) + + # WHEN calling it twice with the same payload + payment = PaymentInput(**MOCK_EVENT) + collect_payment(payment=payment) + second_call = collect_payment(payment=payment) + + # THEN the replayed response is rebuilt as the annotated model + assert isinstance(second_call, PaymentOutput) + + +def test_idempotent_function_dataclass_serializer_deduced_from_postponed_annotation(): + # GIVEN a function annotated with a dataclass under postponed annotations + @idempotent_function( + data_keyword_argument="payment", + persistence_store=NoOpPersistenceLayer(), + config=IdempotencyConfig(use_local_cache=True), + output_serializer=DataclassSerializer, + ) + def collect_payment(payment: dict) -> PaymentOutputDataclass: + return PaymentOutputDataclass(**payment) + + # WHEN calling it twice with the same payload + collect_payment(payment=MOCK_EVENT) + second_call = collect_payment(payment=MOCK_EVENT) + + # THEN the replayed response is rebuilt as the annotated dataclass + assert isinstance(second_call, PaymentOutputDataclass) + assert second_call == PaymentOutputDataclass(**MOCK_EVENT) + + +def test_idempotent_function_pydantic_serializer_postponed_list_annotation_raises_model_type_error(): + idempotent_function_decorator = idempotent_function( + data_keyword_argument="payment", + persistence_store=NoOpPersistenceLayer(), + output_serializer=PydanticSerializer, + ) + + # GIVEN a function annotated with a list of models, which isn't a Pydantic model + # WHEN decorating it + # THEN we raise IdempotencyModelTypeError instead of a TypeError from issubclass + with pytest.raises(IdempotencyModelTypeError, match="Model type is not inherited from pydantic BaseModel"): + + @idempotent_function_decorator + def collect_payment(payment: PaymentInput) -> list[PaymentOutput]: + return [PaymentOutput(**payment.model_dump())] + + +def test_idempotent_function_pydantic_serializer_unresolvable_postponed_annotation_raises_model_type_error(): + # GIVEN a model defined inside a function, which postponed annotations can't resolve from module globals + class LocalPaymentOutput(BaseModel): + customer_id: str + + idempotent_function_decorator = idempotent_function( + data_keyword_argument="payment", + persistence_store=NoOpPersistenceLayer(), + output_serializer=PydanticSerializer, + ) + + # WHEN decorating a function annotated with it + # THEN we raise IdempotencyModelTypeError instead of a TypeError from issubclass + with pytest.raises(IdempotencyModelTypeError, match="Model type is not inherited from pydantic BaseModel"): + + @idempotent_function_decorator + def collect_payment(payment: PaymentInput) -> LocalPaymentOutput: + return LocalPaymentOutput(customer_id=payment.customer_id)