utils.py 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283
  1. """Utility functions for pydantic-settings sources."""
  2. from __future__ import annotations as _annotations
  3. from collections import deque
  4. from collections.abc import Mapping, Sequence
  5. from dataclasses import is_dataclass
  6. from enum import Enum
  7. from typing import Any, TypeVar, cast, get_args, get_origin
  8. from pydantic import BaseModel, Json, RootModel, Secret
  9. from pydantic._internal._utils import is_model_class
  10. from pydantic.dataclasses import is_pydantic_dataclass
  11. from pydantic.fields import FieldInfo
  12. from typing_inspection import typing_objects
  13. from ..exceptions import SettingsError
  14. from ..utils import _lenient_issubclass
  15. from .types import EnvNoneType
  16. def _get_env_var_key(key: str, case_sensitive: bool = False) -> str:
  17. return key if case_sensitive else key.lower()
  18. def _parse_env_none_str(value: str | None, parse_none_str: str | None = None) -> str | None | EnvNoneType:
  19. return value if not (value == parse_none_str and parse_none_str is not None) else EnvNoneType(value)
  20. def parse_env_vars(
  21. env_vars: Mapping[str, str | None],
  22. case_sensitive: bool = False,
  23. ignore_empty: bool = False,
  24. parse_none_str: str | None = None,
  25. ) -> Mapping[str, str | None]:
  26. return {
  27. _get_env_var_key(k, case_sensitive): _parse_env_none_str(v, parse_none_str)
  28. for k, v in env_vars.items()
  29. if not (ignore_empty and v == '')
  30. }
  31. def _substitute_typevars(tp: Any, param_map: dict[Any, Any]) -> Any:
  32. """Substitute TypeVars in a type annotation with concrete types from param_map."""
  33. if isinstance(tp, TypeVar) and tp in param_map:
  34. return param_map[tp]
  35. args = get_args(tp)
  36. if not args:
  37. return tp
  38. new_args = tuple(_substitute_typevars(arg, param_map) for arg in args)
  39. if new_args == args:
  40. return tp
  41. origin = get_origin(tp)
  42. if origin is not None:
  43. try:
  44. return origin[new_args]
  45. except TypeError:
  46. # types.UnionType and similar are not directly subscriptable,
  47. # reconstruct using | operator
  48. import functools
  49. import operator
  50. return functools.reduce(operator.or_, new_args)
  51. return tp
  52. def _resolve_type_alias(annotation: Any) -> Any:
  53. """Resolve a TypeAliasType to its underlying value, substituting type params if parameterized."""
  54. if typing_objects.is_typealiastype(annotation):
  55. return annotation.__value__
  56. origin = get_origin(annotation)
  57. if typing_objects.is_typealiastype(origin):
  58. type_params = getattr(origin, '__type_params__', ())
  59. type_args = get_args(annotation)
  60. value = origin.__value__
  61. if type_params and type_args:
  62. return _substitute_typevars(value, dict(zip(type_params, type_args)))
  63. return value
  64. return annotation
  65. def _annotation_is_complex(annotation: Any, metadata: list[Any]) -> bool:
  66. # If the model is a root model, the root annotation should be used to
  67. # evaluate the complexity.
  68. annotation = _resolve_type_alias(annotation)
  69. if annotation is not None and _lenient_issubclass(annotation, RootModel) and annotation is not RootModel:
  70. annotation = cast('type[RootModel[Any]]', annotation)
  71. root_annotation = annotation.model_fields['root'].annotation
  72. if root_annotation is not None: # pragma: no branch
  73. annotation = root_annotation
  74. if any(isinstance(md, Json) for md in metadata): # type: ignore[misc]
  75. return False
  76. origin = get_origin(annotation)
  77. # Check if annotation is of the form Annotated[type, metadata].
  78. if typing_objects.is_annotated(origin):
  79. # Return result of recursive call on inner type.
  80. inner, *meta = get_args(annotation)
  81. return _annotation_is_complex(inner, meta)
  82. if origin is Secret:
  83. return False
  84. return (
  85. _annotation_is_complex_inner(annotation)
  86. or _annotation_is_complex_inner(origin)
  87. or hasattr(origin, '__pydantic_core_schema__')
  88. or hasattr(origin, '__get_pydantic_core_schema__')
  89. )
  90. def _get_field_metadata(field: FieldInfo) -> list[Any]:
  91. annotation = _resolve_type_alias(field.annotation)
  92. metadata = field.metadata
  93. origin = get_origin(annotation)
  94. if typing_objects.is_annotated(origin):
  95. _, *meta = get_args(annotation)
  96. metadata += meta
  97. return metadata
  98. def _annotation_is_complex_inner(annotation: type[Any] | None) -> bool:
  99. if _lenient_issubclass(annotation, (str, bytes)):
  100. return False
  101. return _lenient_issubclass(
  102. annotation, (BaseModel, Mapping, Sequence, tuple, set, frozenset, deque)
  103. ) or is_dataclass(annotation)
  104. def _union_is_complex(annotation: type[Any] | None, metadata: list[Any]) -> bool:
  105. """Check if a union type contains any complex types."""
  106. return any(_annotation_is_complex(arg, metadata) for arg in get_args(annotation))
  107. def _annotation_contains_types(
  108. annotation: type[Any] | None,
  109. types: tuple[Any, ...],
  110. is_include_origin: bool = True,
  111. is_strip_annotated: bool = False,
  112. is_instance: bool = False,
  113. collect: set[Any] | None = None,
  114. ) -> bool:
  115. """Check if a type annotation contains any of the specified types."""
  116. if is_strip_annotated:
  117. annotation = _strip_annotated(annotation)
  118. if is_include_origin is True:
  119. origin = get_origin(annotation)
  120. if origin in types:
  121. if collect is None:
  122. return True
  123. collect.add(annotation)
  124. if is_instance and any(isinstance(origin, type_) for type_ in types):
  125. if collect is None:
  126. return True
  127. collect.add(annotation)
  128. for type_ in get_args(annotation):
  129. if (
  130. _annotation_contains_types(
  131. type_,
  132. types,
  133. is_include_origin=True,
  134. is_strip_annotated=is_strip_annotated,
  135. is_instance=is_instance,
  136. collect=collect,
  137. )
  138. and collect is None
  139. ):
  140. return True
  141. if is_instance and any(isinstance(annotation, type_) for type_ in types):
  142. if collect is None:
  143. return True
  144. collect.add(annotation)
  145. if annotation in types:
  146. if collect is not None:
  147. collect.add(annotation)
  148. return True
  149. return False
  150. def _strip_annotated(annotation: Any) -> Any:
  151. if typing_objects.is_annotated(get_origin(annotation)):
  152. return annotation.__origin__
  153. else:
  154. return annotation
  155. def _annotation_enum_val_to_name(annotation: type[Any] | None, value: Any) -> str | None:
  156. for type_ in (annotation, get_origin(annotation), *get_args(annotation)):
  157. if _lenient_issubclass(type_, Enum):
  158. if value in type_.__members__.values():
  159. return type_(value).name
  160. return None
  161. def _annotation_enum_name_to_val(annotation: type[Any] | None, name: Any) -> Any:
  162. for type_ in (annotation, get_origin(annotation), *get_args(annotation)):
  163. if _lenient_issubclass(type_, Enum):
  164. if name in type_.__members__.keys():
  165. return type_[name]
  166. return None
  167. def _get_model_fields(model_cls: type[Any]) -> dict[str, Any]:
  168. """Get fields from a pydantic model or dataclass."""
  169. if is_pydantic_dataclass(model_cls) and hasattr(model_cls, '__pydantic_fields__'):
  170. return model_cls.__pydantic_fields__
  171. if is_model_class(model_cls):
  172. return model_cls.model_fields
  173. raise SettingsError(f'Error: {model_cls.__name__} is not subclass of BaseModel or pydantic.dataclasses.dataclass')
  174. def _get_alias_names(
  175. field_name: str,
  176. field_info: Any,
  177. alias_path_args: dict[str, int | None] | None = None,
  178. case_sensitive: bool = True,
  179. ) -> tuple[tuple[str, ...], bool]:
  180. """Get alias names for a field, handling alias paths and case sensitivity."""
  181. from pydantic import AliasChoices, AliasPath
  182. alias_names: list[str] = []
  183. is_alias_path_only: bool = True
  184. if not any((field_info.alias, field_info.validation_alias)):
  185. alias_names += [field_name]
  186. is_alias_path_only = False
  187. else:
  188. new_alias_paths: list[AliasPath] = []
  189. for alias in (field_info.alias, field_info.validation_alias):
  190. if alias is None:
  191. continue
  192. elif isinstance(alias, str):
  193. alias_names.append(alias)
  194. is_alias_path_only = False
  195. elif isinstance(alias, AliasChoices):
  196. for name in alias.choices:
  197. if isinstance(name, str):
  198. alias_names.append(name)
  199. is_alias_path_only = False
  200. else:
  201. new_alias_paths.append(name)
  202. else:
  203. new_alias_paths.append(alias)
  204. for alias_path in new_alias_paths:
  205. name = cast(str, alias_path.path[0])
  206. name = name.lower() if not case_sensitive else name
  207. if alias_path_args is not None:
  208. alias_path_args[name] = (
  209. alias_path.path[1] if len(alias_path.path) > 1 and isinstance(alias_path.path[1], int) else None
  210. )
  211. if not alias_names and is_alias_path_only:
  212. alias_names.append(name)
  213. if not case_sensitive:
  214. alias_names = [alias_name.lower() for alias_name in alias_names]
  215. return tuple(dict.fromkeys(alias_names)), is_alias_path_only
  216. def _is_function(obj: Any) -> bool:
  217. """Check if an object is a function."""
  218. from types import BuiltinFunctionType, FunctionType
  219. return isinstance(obj, (FunctionType, BuiltinFunctionType))
  220. __all__ = [
  221. '_annotation_contains_types',
  222. '_annotation_enum_name_to_val',
  223. '_annotation_enum_val_to_name',
  224. '_annotation_is_complex',
  225. '_annotation_is_complex_inner',
  226. '_get_alias_names',
  227. '_get_env_var_key',
  228. '_get_model_fields',
  229. '_is_function',
  230. '_parse_env_none_str',
  231. '_resolve_type_alias',
  232. '_strip_annotated',
  233. '_union_is_complex',
  234. 'parse_env_vars',
  235. ]