fix: pydantic 2

Author: rubenfonseca

aws_lambda_powertools/event_handler/api_gateway.py

@@ -445,7 +445,7 @@ def dependant(self) -> "Dependant":
         if self._dependant is None:
             from aws_lambda_powertools.event_handler.openapi.dependant import get_dependant
 
-            self._dependant = get_dependant(path=self.openapi_path, call=self.func)
+            self._dependant = get_dependant(path=self.openapi_path, call=self.func, responses=self.responses)
 
         return self._dependant
 
@@ -525,15 +525,15 @@ def _get_openapi_path(
 
                         # Case 2.1: the 'content' has a model
                         if "model" in payload:
-                            from aws_lambda_powertools.event_handler.openapi.params import analyze_param
-
-                            return_field = analyze_param(
-                                param_name="return",
-                                annotation=cast(OpenAPIResponseContentModel, payload)["model"],
-                                value=None,
-                                is_path_param=False,
-                                is_response_param=True,
+                            # Find the model in the dependant's extra models
+                            return_field = next(
+                                filter(
+                                    lambda model: model.type_ is cast(OpenAPIResponseContentModel, payload)["model"],
+                                    self.dependant.response_extra_models,
+                                ),
                             )
+                            if not return_field:
+                                raise AssertionError("Model declared in custom responses was not found")
 
                             new_payload = self._openapi_operation_return(
                                 param=return_field,
@@ -2151,6 +2151,9 @@ def _get_fields_from_routes(routes: Sequence[Route]) -> List["ModelField"]:
             if route.dependant.return_param:
                 responses_from_routes.append(route.dependant.return_param)
 
+            if route.dependant.response_extra_models:
+                responses_from_routes.extend(route.dependant.response_extra_models)
+
         flat_models = list(responses_from_routes + request_fields_from_routes + body_fields_from_routes)
         return flat_models
 

aws_lambda_powertools/event_handler/openapi/dependant.py

@@ -24,6 +24,7 @@
     create_response_field,
     get_flat_dependant,
 )
+from aws_lambda_powertools.event_handler.openapi.types import OpenAPIResponse, OpenAPIResponseContentModel
 
 """
 This turns the opaque function signature into typed, validated models.
@@ -145,6 +146,7 @@ def get_dependant(
     path: str,
     call: Callable[..., Any],
     name: Optional[str] = None,
+    responses: Optional[Dict[int, OpenAPIResponse]] = None,
 ) -> Dependant:
     """
     Returns a dependant model for a handler function. A dependant model is a model that contains
@@ -158,6 +160,8 @@ def get_dependant(
         The handler function
     name: str, optional
         The name of the handler function
+    responses: List[Dict[int, OpenAPIResponse]], optional
+        The list of extra responses for the handler function
 
     Returns
     -------
@@ -210,6 +214,24 @@ def get_dependant(
 
         dependant.return_param = param_field
 
+    # Also add the optional extra responses to the dependant model.
+    if responses:
+        for response in responses.values():
+            if "content" in response:
+                for schema in response["content"].values():
+                    if "model" in schema:
+                        response_field = analyze_param(
+                            param_name="return",
+                            annotation=cast(OpenAPIResponseContentModel, schema)["model"],
+                            value=None,
+                            is_path_param=False,
+                            is_response_param=True,
+                        )
+                        if response_field is None:
+                            raise AssertionError("Response field is None for response model")
+
+                        dependant.response_extra_models.append(response_field)
+
     return dependant
 
 

aws_lambda_powertools/event_handler/openapi/params.py

@@ -49,6 +49,7 @@ def __init__(
         cookie_params: Optional[List[ModelField]] = None,
         body_params: Optional[List[ModelField]] = None,
         return_param: Optional[ModelField] = None,
+        response_extra_models: Optional[List[ModelField]] = None,
         name: Optional[str] = None,
         call: Optional[Callable[..., Any]] = None,
         request_param_name: Optional[str] = None,
@@ -64,6 +65,7 @@ def __init__(
         self.cookie_params = cookie_params or []
         self.body_params = body_params or []
         self.return_param = return_param or None
+        self.response_extra_models = response_extra_models or []
         self.request_param_name = request_param_name
         self.websocket_param_name = websocket_param_name
         self.http_connection_param_name = http_connection_param_name