Add logic to handle globalns and localns in pydantic_model_creator

Author: alistairmaclean

tortoise/contrib/pydantic/creator.py

@@ -61,6 +61,8 @@ def _pydantic_recursion_protector(
     name=None,
     allow_cycles: bool = False,
     sort_alphabetically: Optional[bool] = None,
+    globalns: dict = None,
+    localns: dict = None,
 ) -> Optional[type[PydanticModel]]:
     """
     It is an inner function to protect pydantic model creator against cyclic recursion
@@ -93,6 +95,8 @@ def _pydantic_recursion_protector(
         _stack=stack,
         allow_cycles=allow_cycles,
         sort_alphabetically=sort_alphabetically,
+        globalns=globalns,
+        localns=localns,
         _as_submodel=True,
     )
     return pmc.create_pydantic_model()
@@ -237,6 +241,8 @@ def __init__(
         model_config: Optional[ConfigDict] = None,
         validators: Optional[dict[str, Any]] = None,
         module: str = __name__,
+        globalns: dict = None,
+        localns: dict = None,
         _stack: tuple = (),
         _as_submodel: bool = False,
     ) -> None:
@@ -288,7 +294,7 @@ def __init__(
 
         self._as_submodel = _as_submodel
 
-        self._annotations = get_annotations(cls)
+        self._annotations = get_annotations(cls, globalns=globalns, localns=localns)
 
         self._pconfig: ConfigDict
 
@@ -305,6 +311,9 @@ def __init__(
         self._validators = validators
         self._module = module
 
+        self.globalns = globalns
+        self.localns = localns
+
         self._stack = _stack
 
     @property
@@ -523,7 +532,7 @@ def _process_computed_field(
         field: ComputedFieldDescription,
     ) -> Optional[Any]:
         func = field.function
-        annotation = get_annotations(self._cls, func).get("return", None)
+        annotation = get_annotations(self._cls, func, globalns=self.globalns, localns=self.localns).get("return", None)
         comment = _cleandoc(func)
         if annotation is not None:
             c_f = computed_field(return_type=annotation, description=comment)
@@ -555,6 +564,8 @@ def get_fields_to_carry_on(field_tuple: tuple[str, ...]) -> tuple[str, ...]:
                 stack=new_stack,
                 allow_cycles=self.meta.allow_cycles,
                 sort_alphabetically=self.meta.sort_alphabetically,
+                globalns=self.globalns,
+                localns=self.localns,
             )
         else:
             pmodel = None
@@ -581,6 +592,8 @@ def pydantic_model_creator(
     model_config: Optional[ConfigDict] = None,
     validators: Optional[dict[str, Any]] = None,
     module: str = __name__,
+    globalns: dict = None,
+    localns: dict = None,
 ) -> type[PydanticModel]:
     """
     Function to build `Pydantic Model <https://docs.pydantic.dev/latest/concepts/models/>`__ off Tortoise Model.
@@ -607,6 +620,8 @@ def pydantic_model_creator(
     :param model_config: A custom config to use as pydantic config.
     :param validators: A dictionary of methods that validate fields.
     :param module: The name of the module that the model belongs to.
+    :param globalns: If specified, use this dictionary as the globals map.
+    :param localns: If specified, use this dictionary as the locals map.
 
         Note: Created pydantic model uses config_class parameter and PydanticMeta's
             config_class as its Config class's bases(Only if provided!), but it
@@ -627,5 +642,7 @@ def pydantic_model_creator(
         model_config=model_config,
         validators=validators,
         module=module,
+        globalns=globalns,
+        localns=localns,
     )
     return pmc.create_pydantic_model()

tortoise/contrib/pydantic/utils.py

@@ -5,11 +5,21 @@
     from tortoise.models import Model
 
 
+<<<<<<< HEAD
 def get_annotations(cls: "type[Model]", method: Optional[Callable] = None) -> dict[str, Any]:
+=======
+def get_annotations(cls: "Type[Model]", method: Optional[Callable] = None, globalns: dict = None, localns: dict = None) -> dict[str, Any]:
+>>>>>>> Add logic to handle `globalns` and `localns` in `pydantic_model_creator`
     """
     Get all annotations including base classes
     :param cls: The model class we need annotations from
     :param method: If specified, we try to get the annotations for the callable
+    :param globalns: If specified, use this dictionary as the globals map
+    :param localns: If specified, use this dictionary as the locals map
     :return: The list of annotations
     """
+<<<<<<< HEAD
     return get_type_hints(method or cls)
+=======
+    return typing.get_type_hints(method or cls, globalns, localns)
+>>>>>>> Add logic to handle `globalns` and `localns` in `pydantic_model_creator`