Skip to content

Commit

Permalink
fix(parser): support TypeAdapter instances as models (#5535)
Browse files Browse the repository at this point in the history
  • Loading branch information
leandrodamascena authored Nov 11, 2024
1 parent 68ddf3e commit 217f985
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 2 deletions.
7 changes: 6 additions & 1 deletion aws_lambda_powertools/utilities/parser/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,17 @@ def _retrieve_or_set_model_from_cache(model: type[T]) -> TypeAdapter:
The TypeAdapter instance for the given model,
either retrieved from the cache or newly created and stored in the cache.
"""

id_model = id(model)

if id_model in CACHE_TYPE_ADAPTER:
return CACHE_TYPE_ADAPTER[id_model]

CACHE_TYPE_ADAPTER[id_model] = TypeAdapter(model)
if isinstance(model, TypeAdapter):
CACHE_TYPE_ADAPTER[id_model] = model
else:
CACHE_TYPE_ADAPTER[id_model] = TypeAdapter(model)

return CACHE_TYPE_ADAPTER[id_model]


Expand Down
36 changes: 35 additions & 1 deletion tests/functional/parser/test_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,41 @@ class FailedCallback(pydantic.BaseModel):
DogCallback = Annotated[Union[SuccessfulCallback, FailedCallback], pydantic.Field(discriminator="status")]

@event_parser(model=DogCallback)
def handler(event: test_input, _: Any) -> str:
def handler(event, _: Any) -> str:
if isinstance(event, FailedCallback):
return f"Uh oh. Had a problem: {event.error}"

return f"Successfully retrieved {event.breed} named {event.name}"

ret = handler(test_input, None)
assert ret == expected


@pytest.mark.parametrize(
"test_input,expected",
[
(
{"status": "succeeded", "name": "Clifford", "breed": "Labrador"},
"Successfully retrieved Labrador named Clifford",
),
({"status": "failed", "error": "oh some error"}, "Uh oh. Had a problem: oh some error"),
],
)
def test_parser_unions_with_type_adapter_instance(test_input, expected):
class SuccessfulCallback(pydantic.BaseModel):
status: Literal["succeeded"]
name: str
breed: Literal["Newfoundland", "Labrador"]

class FailedCallback(pydantic.BaseModel):
status: Literal["failed"]
error: str

DogCallback = Annotated[Union[SuccessfulCallback, FailedCallback], pydantic.Field(discriminator="status")]
DogCallbackTypeAdapter = pydantic.TypeAdapter(DogCallback)

@event_parser(model=DogCallbackTypeAdapter)
def handler(event, _: Any) -> str:
if isinstance(event, FailedCallback):
return f"Uh oh. Had a problem: {event.error}"

Expand Down

0 comments on commit 217f985

Please sign in to comment.