fix: tests

Author: rubenfonseca

aws_lambda_powertools/event_handler/api_gateway.py

@@ -263,6 +263,9 @@ def __init__(
         # _middleware_stack_built is used to ensure the middleware stack is only built once.
         self._middleware_stack_built = False
 
+        # _dependant is used to cache the dependant model for the handler function
+        self._dependant: Optional["Dependant"] = None
+
     def __call__(
         self,
         router_middlewares: List[Callable],
@@ -353,6 +356,15 @@ def _build_middleware_stack(self, router_middlewares: List[Callable[..., Any]])
 
         self._middleware_stack_built = True
 
+    @property
+    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.path, call=self.func)
+
+        return self._dependant
+
     def _get_openapi_path(
         self,
         *,
@@ -1168,8 +1180,7 @@ def get_openapi_schema(
             get_compat_model_name_map,
             get_definitions,
         )
-        from aws_lambda_powertools.event_handler.openapi.dependant import get_dependant
-        from aws_lambda_powertools.event_handler.openapi.models import OpenAPI, Server
+        from aws_lambda_powertools.event_handler.openapi.models import OpenAPI, PathItem, Server
         from aws_lambda_powertools.event_handler.openapi.types import (
             COMPONENT_REF_TEMPLATE,
         )
@@ -1213,13 +1224,8 @@ def get_openapi_schema(
 
         # Add routes to the OpenAPI schema
         for route in all_routes:
-            dependant = get_dependant(
-                path=route.path,
-                call=route.func,
-            )
-
             result = route._get_openapi_path(
-                dependant=dependant,
+                dependant=route.dependant,
                 operation_ids=operation_ids,
                 model_name_map=model_name_map,
                 field_mapping=field_mapping,
@@ -1238,7 +1244,7 @@ def get_openapi_schema(
         if tags:
             output["tags"] = tags
 
-        output["paths"] = paths
+        output["paths"] = {k: PathItem(**v) for k, v in paths.items()}
 
         return OpenAPI(**output)
 
@@ -1701,20 +1707,18 @@ def _get_fields_from_routes(routes: Sequence[Route]) -> List["ModelField"]:
         """
 
         from aws_lambda_powertools.event_handler.openapi.dependant import (
-            get_dependant,
             get_flat_params,
         )
 
         responses_from_routes: List["ModelField"] = []
         request_fields_from_routes: List["ModelField"] = []
 
         for route in routes:
-            dependant = get_dependant(path=route.path, call=route.func)
-            params = get_flat_params(dependant)
+            params = get_flat_params(route.dependant)
             request_fields_from_routes.extend(params)
 
-            if dependant.return_param:
-                responses_from_routes.append(dependant.return_param)
+            if route.dependant.return_param:
+                responses_from_routes.append(route.dependant.return_param)
 
         flat_models = list(responses_from_routes + request_fields_from_routes)
         return flat_models

aws_lambda_powertools/event_handler/openapi/compat.py

@@ -1,5 +1,6 @@
 # mypy: ignore-errors
 # flake8: noqa
+from copy import copy
 
 # MAINTENANCE: remove when deprecating Pydantic v1. Mypy doesn't handle two different code paths that import different
 # versions of a module, so we need to ignore errors here.
@@ -10,19 +11,27 @@
 
 from typing_extensions import Annotated, Literal
 
-from aws_lambda_powertools.event_handler.openapi.types import COMPONENT_REF_PREFIX, PYDANTIC_V2, ModelNameMap
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+
+from aws_lambda_powertools.event_handler.openapi.types import (
+    COMPONENT_REF_PREFIX,
+    PYDANTIC_V2,
+    ModelNameMap,
+)
 
 if PYDANTIC_V2:
     from pydantic import TypeAdapter
     from pydantic._internal._typing_extra import eval_type_lenient
     from pydantic.fields import FieldInfo
     from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
-    from pydantic_core import PydanticUndefined
+    from pydantic_core import PydanticUndefined, PydanticUndefinedType
 
     from aws_lambda_powertools.event_handler.openapi.types import IncEx
 
     Undefined = PydanticUndefined
     Required = PydanticUndefined
+    UndefinedType = PydanticUndefinedType
 
     evaluate_forwardref = eval_type_lenient
 
@@ -49,7 +58,7 @@ def default(self) -> Any:
         def type_(self) -> Any:
             return self.field_info.annotation
 
-        def __post__init__(self) -> None:
+        def __post_init__(self) -> None:
             self._type_adapter: TypeAdapter[Any] = TypeAdapter(
                 Annotated[self.field_info.annotation, self.field_info],
             )
@@ -125,9 +134,15 @@ def get_compat_model_name_map(fields: List[ModelField]) -> ModelNameMap:
     def get_annotation_from_field_info(annotation: Any, field_info: FieldInfo, field_name: str) -> Any:
         return annotation
 
+    def model_rebuild(model: Type[BaseModel]) -> None:
+        model.model_rebuild()
+
+    def copy_field_info(*, field_info: FieldInfo, annotation: Any) -> FieldInfo:
+        return type(field_info).from_annotation(annotation)
+
 else:
     from pydantic import BaseModel
-    from pydantic.fields import ModelField, Required, Undefined
+    from pydantic.fields import ModelField, Required, Undefined, UndefinedType
     from pydantic.schema import (
         field_schema,
         get_annotation_from_field_info,
@@ -192,3 +207,9 @@ def get_model_definitions(
     def get_compat_model_name_map(fields: List[ModelField]) -> ModelNameMap:
         models = get_flat_models_from_fields(fields, known_models=set())
         return get_model_name_map(models)
+
+    def model_rebuild(model: Type[BaseModel]) -> None:
+        model.update_forward_refs()
+
+    def copy_field_info(*, field_info: FieldInfo, annotation: Any) -> FieldInfo:
+        return copy(field_info)

aws_lambda_powertools/event_handler/openapi/dependant.py

@@ -160,11 +160,12 @@ def get_dependant(
         is_path_param = param_name in path_param_names
 
         # Analyze the parameter to get the Pydantic field.
-        _, param_field = analyze_param(
+        param_field = analyze_param(
             param_name=param_name,
             annotation=param.annotation,
             value=param.default,
             is_path_param=is_path_param,
+            is_response_param=False,
         )
         if param_field is None:
             raise AssertionError(f"Param field is None for param: {param_name}")
@@ -174,11 +175,12 @@ def get_dependant(
     # If the return annotation is not empty, add it to the dependant model.
     return_annotation = endpoint_signature.return_annotation
     if return_annotation is not inspect.Signature.empty:
-        _, param_field = analyze_param(
-            param_name="Return",
+        param_field = analyze_param(
+            param_name="return",
             annotation=return_annotation,
             value=None,
             is_path_param=False,
+            is_response_param=True,
         )
         if param_field is None:
             raise AssertionError("Param field is None for return annotation")

aws_lambda_powertools/event_handler/openapi/models.py

@@ -4,6 +4,7 @@
 from pydantic import AnyUrl, BaseModel, Field
 from typing_extensions import Annotated, Literal
 
+from aws_lambda_powertools.event_handler.openapi.compat import model_rebuild
 from aws_lambda_powertools.event_handler.openapi.types import PYDANTIC_V2
 
 """
@@ -205,7 +206,7 @@ class Schema(BaseModel):
     deprecated: Optional[bool] = None
     readOnly: Optional[bool] = None
     writeOnly: Optional[bool] = None
-    examples: Optional[List[Any]] = None
+    examples: Optional[List["Example"]] = None
     # Ref: OpenAPI 3.1.0: https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.1.0.md#schema-object
     # Schema Object
     discriminator: Optional[Discriminator] = None
@@ -577,6 +578,6 @@ class Config:
             extra = "allow"
 
 
-Schema.update_forward_refs()
-Operation.update_forward_refs()
-Encoding.update_forward_refs()
+model_rebuild(Schema)
+model_rebuild(Operation)
+model_rebuild(Encoding)

aws_lambda_powertools/event_handler/openapi/params.py

@@ -1,16 +1,17 @@
 import inspect
-from copy import copy
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 from pydantic import BaseConfig
 from pydantic.fields import FieldInfo
-from typing_extensions import Annotated, get_args, get_origin
+from typing_extensions import Annotated, Literal, get_args, get_origin
 
 from aws_lambda_powertools.event_handler.openapi.compat import (
     ModelField,
     Required,
     Undefined,
+    UndefinedType,
+    copy_field_info,
     get_annotation_from_field_info,
 )
 from aws_lambda_powertools.event_handler.openapi.types import PYDANTIC_V2, CacheKey
@@ -302,7 +303,8 @@ def analyze_param(
     annotation: Any,
     value: Any,
     is_path_param: bool,
-) -> Tuple[Any, Optional[ModelField]]:
+    is_response_param: bool,
+) -> Optional[ModelField]:
     """
     Analyze a parameter annotation and value to determine the type and default value of the parameter.
 
@@ -316,10 +318,12 @@ def analyze_param(
         The value of the parameter
     is_path_param
         Whether the parameter is a path parameter
+    is_response_param
+        Whether the parameter is the return annotation
 
     Returns
     -------
-    Tuple[Any, Optional[ModelField]]
+    Optional[ModelField]
         The type annotation and the Pydantic field representing the parameter
     """
     field_info, type_annotation = _get_field_info_and_type_annotation(annotation, value, is_path_param)
@@ -336,12 +340,16 @@ def analyze_param(
 
         # Check if the parameter is part of the path. Otherwise, defaults to query.
         if is_path_param:
-            field_info = Path(annotation=type_annotation, default=default_value)
+            field_info = Path(annotation=type_annotation)
         else:
             field_info = Query(annotation=type_annotation, default=default_value)
 
+    # When we have a response field, we need to set the default value to Required
+    if is_response_param:
+        field_info.default = Required
+
     field = _create_model_field(field_info, type_annotation, param_name, is_path_param)
-    return type_annotation, field
+    return field
 
 
 def _get_field_info_and_type_annotation(annotation, value, is_path_param: bool) -> Tuple[Optional[FieldInfo], Any]:
@@ -372,7 +380,10 @@ def _get_field_info_annotated_type(annotation, value, is_path_param: bool) -> Tu
 
     if isinstance(powertools_annotation, FieldInfo):
         # Copy `field_info` because we mutate `field_info.default` later
-        field_info = copy(powertools_annotation)
+        field_info = copy_field_info(
+            field_info=powertools_annotation,
+            annotation=annotation,
+        )
         if field_info.default not in [Undefined, Required]:
             raise AssertionError("FieldInfo needs to have a default value of Undefined or Required")
 
@@ -386,6 +397,44 @@ def _get_field_info_annotated_type(annotation, value, is_path_param: bool) -> Tu
     return field_info, type_annotation
 
 
+def _create_response_field(
+    name: str,
+    type_: Type[Any],
+    default: Optional[Any] = Undefined,
+    required: Union[bool, UndefinedType] = Undefined,
+    model_config: Type[BaseConfig] = BaseConfig,
+    field_info: Optional[FieldInfo] = None,
+    alias: Optional[str] = None,
+    mode: Literal["validation", "serialization"] = "validation",
+) -> ModelField:
+    """
+    Create a new response field. Raises if type_ is invalid.
+    """
+    if PYDANTIC_V2:
+        field_info = field_info or FieldInfo(
+            annotation=type_,
+            default=default,
+            alias=alias,
+        )
+    else:
+        field_info = field_info or FieldInfo()
+    kwargs = {"name": name, "field_info": field_info}
+    if PYDANTIC_V2:
+        kwargs.update({"mode": mode})
+    else:
+        kwargs.update(
+            {
+                "type_": type_,
+                "class_validators": {},
+                "default": default,
+                "required": required,
+                "model_config": model_config,
+                "alias": alias,
+            },
+        )
+    return ModelField(**kwargs)  # type: ignore[arg-type]
+
+
 def _create_model_field(
     field_info: Optional[FieldInfo],
     type_annotation: Any,
@@ -411,21 +460,11 @@ def _create_model_field(
         alias = field_info.alias or param_name
     field_info.alias = alias
 
-    # Create the Pydantic field
-    kwargs = {"name": param_name, "field_info": field_info}
-
-    if PYDANTIC_V2:
-        kwargs.update({"mode": "validation"})
-    else:
-        kwargs.update(
-            {
-                "type_": use_annotation,
-                "class_validators": {},
-                "default": field_info.default,
-                "required": field_info.default in (Required, Undefined),
-                "model_config": BaseConfig,
-                "alias": alias,
-            },
-        )
-
-    return ModelField(**kwargs)  # type: ignore[arg-type]
+    return _create_response_field(
+        name=param_name,
+        type_=use_annotation,
+        default=field_info.default,
+        alias=alias,
+        required=field_info.default in (Required, Undefined),
+        field_info=field_info,
+    )

tests/functional/event_handler/test_openapi_params.py

@@ -98,7 +98,7 @@ def test_openapi_with_custom_params():
     def handler(
         count: Annotated[
             int,
-            Query(lt=100, gt=0, examples=[Example(summary="Example 1", value=10)]),
+            Query(gt=0, lt=100, examples=[Example(summary="Example 1", value=10)]),
         ] = 1,
     ):
         raise NotImplementedError()