diff --git a/aws_lambda_powertools/utilities/data_classes/event_source.py b/aws_lambda_powertools/utilities/data_classes/event_source.py index 6096e3ae7bc..58c8be968f5 100644 --- a/aws_lambda_powertools/utilities/data_classes/event_source.py +++ b/aws_lambda_powertools/utilities/data_classes/event_source.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast from aws_lambda_powertools.middleware_factory import lambda_handler_decorator @@ -10,14 +10,30 @@ from aws_lambda_powertools.utilities.data_classes.common import DictWrapper from aws_lambda_powertools.utilities.typing import LambdaContext +DataClassT = TypeVar("DataClassT", bound="DictWrapper") +OutputT = TypeVar("OutputT") + + +class _EventSourceDecorator(Protocol): + """Annotation of event_source, that lambda_handler_decorator erases at runtime.""" + + def __call__( + self, + *, + data_class: type[DataClassT], + ) -> Callable[ + [Callable[[DataClassT, LambdaContext], OutputT]], + Callable[[dict[str, Any], LambdaContext], OutputT], + ]: ... + @lambda_handler_decorator -def event_source( - handler: Callable[[Any, LambdaContext], Any], +def _event_source( + handler: Callable[[Any, LambdaContext], OutputT], event: dict[str, Any], context: LambdaContext, - data_class: type[DictWrapper], -): + data_class: type[DataClassT], +) -> OutputT: """Middleware to create an instance of the passed in event source data class Parameters @@ -43,3 +59,6 @@ def handler(event: S3Event, context): return {"key": event.object_key} """ return handler(data_class(event), context) + + +event_source = cast(_EventSourceDecorator, _event_source) diff --git a/examples/event_sources/src/kinesis_firehose_response_exception.py b/examples/event_sources/src/kinesis_firehose_response_exception.py index 43ba3a039b2..5363107eea8 100644 --- a/examples/event_sources/src/kinesis_firehose_response_exception.py +++ b/examples/event_sources/src/kinesis_firehose_response_exception.py @@ -9,11 +9,10 @@ @event_source(data_class=KinesisFirehoseEvent) -def lambda_handler(event: dict, context: LambdaContext): - firehose_event = KinesisFirehoseEvent(event) +def lambda_handler(event: KinesisFirehoseEvent, context: LambdaContext): result = KinesisFirehoseDataTransformationResponse() - for record in firehose_event.records: + for record in event.records: try: payload = record.data_as_text # base64 decoded data as str