Skip to content

Commit 217f985

Browse files
fix(parser): support TypeAdapter instances as models (#5535)
1 parent 68ddf3e commit 217f985

File tree

2 files changed

+41
-2
lines changed

2 files changed

+41
-2
lines changed

aws_lambda_powertools/utilities/parser/functions.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,17 @@ def _retrieve_or_set_model_from_cache(model: type[T]) -> TypeAdapter:
3535
The TypeAdapter instance for the given model,
3636
either retrieved from the cache or newly created and stored in the cache.
3737
"""
38+
3839
id_model = id(model)
3940

4041
if id_model in CACHE_TYPE_ADAPTER:
4142
return CACHE_TYPE_ADAPTER[id_model]
4243

43-
CACHE_TYPE_ADAPTER[id_model] = TypeAdapter(model)
44+
if isinstance(model, TypeAdapter):
45+
CACHE_TYPE_ADAPTER[id_model] = model
46+
else:
47+
CACHE_TYPE_ADAPTER[id_model] = TypeAdapter(model)
48+
4449
return CACHE_TYPE_ADAPTER[id_model]
4550

4651

tests/functional/parser/test_parser.py

+35-1
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,41 @@ class FailedCallback(pydantic.BaseModel):
154154
DogCallback = Annotated[Union[SuccessfulCallback, FailedCallback], pydantic.Field(discriminator="status")]
155155

156156
@event_parser(model=DogCallback)
157-
def handler(event: test_input, _: Any) -> str:
157+
def handler(event, _: Any) -> str:
158+
if isinstance(event, FailedCallback):
159+
return f"Uh oh. Had a problem: {event.error}"
160+
161+
return f"Successfully retrieved {event.breed} named {event.name}"
162+
163+
ret = handler(test_input, None)
164+
assert ret == expected
165+
166+
167+
@pytest.mark.parametrize(
168+
"test_input,expected",
169+
[
170+
(
171+
{"status": "succeeded", "name": "Clifford", "breed": "Labrador"},
172+
"Successfully retrieved Labrador named Clifford",
173+
),
174+
({"status": "failed", "error": "oh some error"}, "Uh oh. Had a problem: oh some error"),
175+
],
176+
)
177+
def test_parser_unions_with_type_adapter_instance(test_input, expected):
178+
class SuccessfulCallback(pydantic.BaseModel):
179+
status: Literal["succeeded"]
180+
name: str
181+
breed: Literal["Newfoundland", "Labrador"]
182+
183+
class FailedCallback(pydantic.BaseModel):
184+
status: Literal["failed"]
185+
error: str
186+
187+
DogCallback = Annotated[Union[SuccessfulCallback, FailedCallback], pydantic.Field(discriminator="status")]
188+
DogCallbackTypeAdapter = pydantic.TypeAdapter(DogCallback)
189+
190+
@event_parser(model=DogCallbackTypeAdapter)
191+
def handler(event, _: Any) -> str:
158192
if isinstance(event, FailedCallback):
159193
return f"Uh oh. Had a problem: {event.error}"
160194

0 commit comments

Comments
 (0)