From bf25bfd88eec4351487a40f6b8f102de37130e7f Mon Sep 17 00:00:00 2001 From: Quoc Duong Date: Tue, 17 Sep 2024 11:26:27 +0700 Subject: [PATCH] Apply appropriate method for function wrapping --- src/extendable_pydantic/_patch.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/extendable_pydantic/_patch.py b/src/extendable_pydantic/_patch.py index 561873c..f0cc67d 100644 --- a/src/extendable_pydantic/_patch.py +++ b/src/extendable_pydantic/_patch.py @@ -66,12 +66,13 @@ def _create_response_field_wrapper(wrapped, instance, args, kwargs): _resolve_model_fields_annotation([field]) return field - # fastapi < 0.112.3 - wrapt.wrap_function_wrapper( - utils, "create_response_field", _create_response_field_wrapper - ) - - # fastapi >= 0.112.3 - wrapt.wrap_function_wrapper( - utils, "create_model_field", _create_response_field_wrapper - ) + if hasattr(utils, "create_response_field"): + # For fastapi < 0.112.3 + wrapt.wrap_function_wrapper( + utils, "create_response_field", _create_response_field_wrapper + ) + else: + # For fastapi >= 0.112.3 + wrapt.wrap_function_wrapper( + utils, "create_model_field", _create_response_field_wrapper + )