fix(idempotency): resolve postponed return annotations for serializer classes

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()
 

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)

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"
         ```
 

examples/idempotency/src/working_with_dataclass_deduced_output_serializer.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 from dataclasses import dataclass
 

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 (

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",
+            "<external_wrapper>",
+            "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",
+            "<native_deferred_annotations>",
+            "exec",
+            dont_inherit=True,
+        ),
+        namespace,
+    )
+    replay(namespace["collect_payment"], PydanticSerializer, PaymentOutput(**MOCK_EVENT))
+    assert len(calls) == 1