Apply appropriate method for function wrapping

lmignon · Sep 17, 2024 · a8b28e1 · a8b28e1
1 parent ffdc03a
commit a8b28e1
Showing 1 changed file with 10 additions and 9 deletions.
diff --git 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
+        )
+    if hasattr(utils, "create_model_field"):
+        # For fastapi >= 0.112.3
+        wrapt.wrap_function_wrapper(
+            utils, "create_model_field", _create_response_field_wrapper
+        )