fix: use generics

Author: rubenfonseca

aws_lambda_powertools/event_handler/api_gateway.py

@@ -15,6 +15,7 @@
     Any,
     Callable,
     Dict,
+    Generic,
     List,
     Match,
     Optional,
@@ -23,6 +24,7 @@
     Set,
     Tuple,
     Type,
+    TypeVar,
     Union,
     cast,
 )
@@ -63,6 +65,8 @@
 _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION = "Successful Response"
 _ROUTE_REGEX = "^{}$"
 
+ResponseEventT = TypeVar("ResponseEventT", bound=BaseProxyEvent)
+
 if TYPE_CHECKING:
     from aws_lambda_powertools.event_handler.openapi.compat import (
         JsonSchemaValue,
@@ -691,14 +695,14 @@ def _generate_operation_id(self) -> str:
         return operation_id
 
 
-class ResponseBuilder:
+class ResponseBuilder(Generic[ResponseEventT]):
     """Internally used Response builder"""
 
     def __init__(self, response: Response, route: Optional[Route] = None):
         self.response = response
         self.route = route
 
-    def _add_cors(self, event: BaseProxyEvent, cors: CORSConfig):
+    def _add_cors(self, event: ResponseEventT, cors: CORSConfig):
         """Update headers to include the configured Access-Control headers"""
         self.response.headers.update(cors.to_dict(event.get_header_value("Origin")))
 
@@ -711,7 +715,7 @@ def _add_cache_control(self, cache_control: str):
     def _has_compression_enabled(
         route_compression: bool,
         response_compression: Optional[bool],
-        event: BaseProxyEvent,
+        event: ResponseEventT,
     ) -> bool:
         """
         Checks if compression is enabled.
@@ -724,7 +728,7 @@ def _has_compression_enabled(
             A boolean indicating whether compression is enabled or not in the route setting.
         response_compression: bool, optional
             A boolean indicating whether compression is enabled or not in the response setting.
-        event: BaseProxyEvent
+        event: Generic[ResponseEventT]
             The event object containing the request details.
 
         Returns
@@ -754,7 +758,7 @@ def _compress(self):
         gzip = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS | 16)
         self.response.body = gzip.compress(self.response.body) + gzip.flush()
 
-    def _route(self, event: BaseProxyEvent, cors: Optional[CORSConfig]):
+    def _route(self, event: ResponseEventT, cors: Optional[CORSConfig]):
         """Optionally handle any of the route's configure response handling"""
         if self.route is None:
             return
@@ -769,7 +773,7 @@ def _route(self, event: BaseProxyEvent, cors: Optional[CORSConfig]):
         ):
             self._compress()
 
-    def build(self, event: BaseProxyEvent, cors: Optional[CORSConfig] = None) -> Dict[str, Any]:
+    def build(self, event: ResponseEventT, cors: Optional[CORSConfig] = None) -> Dict[str, Any]:
         """Build the full response dict to be returned by the lambda"""
         self._route(event, cors)
 
@@ -1317,7 +1321,7 @@ def __init__(
         self._strip_prefixes = strip_prefixes
         self.context: Dict = {}  # early init as customers might add context before event resolution
         self.processed_stack_frames = []
-        self._response_builder_class = ResponseBuilder
+        self._response_builder_class = ResponseBuilder[BaseProxyEvent]
 
         # Allow for a custom serializer or a concise json serialization
         self._serializer = serializer or partial(json.dumps, separators=(",", ":"), cls=Encoder)

aws_lambda_powertools/event_handler/bedrock_agent.py

@@ -1,30 +1,27 @@
 import logging
-from typing import Any, Dict, Optional, cast
+from typing import Any, Dict, Optional
 
 from typing_extensions import override
 
 from aws_lambda_powertools.event_handler import ApiGatewayResolver
 from aws_lambda_powertools.event_handler.api_gateway import CORSConfig, ProxyEventType, ResponseBuilder
 from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
-from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent
 
 logger = logging.getLogger(__name__)
 
 
 class BedrockResponseBuilder(ResponseBuilder):
     @override
-    def build(self, event: BaseProxyEvent, cors: Optional[CORSConfig] = None) -> Dict[str, Any]:
+    def build(self, event: BedrockAgentEvent, cors: Optional[CORSConfig] = None) -> Dict[str, Any]:
         """Build the full response dict to be returned by the lambda"""
         self._route(event, cors)
 
-        bedrock_event = cast(BedrockAgentEvent, event)
-
         return {
             "messageVersion": "1.0",
             "response": {
-                "actionGroup": bedrock_event.action_group,
-                "apiPath": bedrock_event.api_path,
-                "httpMethod": bedrock_event.http_method,
+                "actionGroup": event.action_group,
+                "apiPath": event.api_path,
+                "httpMethod": event.http_method,
                 "httpStatusCode": self.response.status_code,
                 "responseBody": {
                     "application/json": {