Use Protocol for type checking

Author: ConnorKirk

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -28,7 +28,7 @@
     from aws_lambda_powertools.utilities.idempotency.persistence.base import (
         BasePersistenceLayer,
     )
-    from aws_lambda_powertools.utilities.typing import DurableContext, LambdaContext
+    from aws_lambda_powertools.utilities.typing import DurableContextProtocol, LambdaContext
 
 from aws_lambda_powertools.warnings import PowertoolsUserWarning
 
@@ -37,9 +37,9 @@
 
 @lambda_handler_decorator
 def idempotent(
-    handler: Callable[[Any, LambdaContext | DurableContext], Any],
+    handler: Callable[[Any, LambdaContext | DurableContextProtocol], Any],
     event: dict[str, Any],
-    context: LambdaContext | DurableContext,
+    context: LambdaContext | DurableContextProtocol,
     persistence_store: BasePersistenceLayer,
     config: IdempotencyConfig | None = None,
     key_prefix: str | None = None,
@@ -92,9 +92,9 @@ def handler(event, context):
 
     config = config or IdempotencyConfig()
 
-    if hasattr(context, "state"):
+    if hasattr(context, "state") and hasattr(context, "lambda_context"):
         # Extract lambda_context from DurableContext
-        durable_context = cast("DurableContext", context)
+        durable_context = cast("DurableContextProtocol", context)
         config.register_lambda_context(durable_context.lambda_context)
         # Note: state.operations is accessed via duck typing at runtime
         is_replay = len(durable_context.state.operations) > 1  # type: ignore[attr-defined]

aws_lambda_powertools/utilities/typing/__init__.py

@@ -4,6 +4,6 @@
     [`Typing`](../utilities/typing.md)
 """
 
-from .lambda_context import DurableContext, LambdaContext
+from .lambda_context import DurableContextProtocol, LambdaContext
 
-__all__ = ["DurableContext", "LambdaContext"]
+__all__ = ["DurableContextProtocol", "LambdaContext"]

aws_lambda_powertools/utilities/typing/lambda_context.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
     from aws_lambda_powertools.utilities.typing.lambda_client_context import (
@@ -95,14 +95,13 @@ def get_remaining_time_in_millis() -> int:
         return 0
 
 
-class DurableContext:
+@runtime_checkable
+class DurableContextProtocol(Protocol):
     _lambda_context: LambdaContext
     _state: object
 
     @property
-    def lambda_context(self) -> LambdaContext:
-        return self._lambda_context
+    def lambda_context(self) -> LambdaContext: ...
 
     @property
-    def state(self) -> object:
-        return self._state
+    def state(self) -> object: ...

tests/functional/idempotency/_boto3/test_idempotency.py

@@ -46,7 +46,7 @@
 from aws_lambda_powertools.utilities.idempotency.serialization.dataclass import (
     DataclassSerializer,
 )
-from aws_lambda_powertools.utilities.typing import DurableContext
+from aws_lambda_powertools.utilities.typing import DurableContextProtocol
 from aws_lambda_powertools.utilities.validation import envelopes, validator
 from aws_lambda_powertools.warnings import PowertoolsUserWarning
 from tests.functional.idempotency.utils import (
@@ -2145,18 +2145,18 @@ def lambda_handler(record, context):
 @pytest.fixture
 def durable_context_single_operation(lambda_context):
     """DurableContext with single operation (execution mode, is_replay=False)"""
-    durable_ctx = DurableContext()
-    durable_ctx._lambda_context = lambda_context
-    durable_ctx._state = Mock(operations=[{"id": "op1"}])
+    durable_ctx = Mock(spec=DurableContextProtocol)
+    durable_ctx.lambda_context = lambda_context
+    durable_ctx.state = Mock(operations=[{"id": "op1"}])
     return durable_ctx
 
 
 @pytest.fixture
 def durable_context_multiple_operations(lambda_context):
     """DurableContext with multiple operations (replay mode, is_replay=True)"""
-    durable_ctx = DurableContext()
-    durable_ctx._lambda_context = lambda_context
-    durable_ctx._state = Mock(operations=[{"id": "op1"}, {"id": "op2"}])
+    durable_ctx = Mock(spec=DurableContextProtocol)
+    durable_ctx.lambda_context = lambda_context
+    durable_ctx.state = Mock(operations=[{"id": "op1"}, {"id": "op2"}])
     return durable_ctx