sources.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649
  1. from __future__ import annotations as _annotations
  2. import json
  3. import os
  4. import warnings
  5. from abc import ABC, abstractmethod
  6. from collections import deque
  7. from dataclasses import is_dataclass
  8. from pathlib import Path
  9. from typing import TYPE_CHECKING, Any, List, Mapping, Sequence, Tuple, Union, cast
  10. from dotenv import dotenv_values
  11. from pydantic import AliasChoices, AliasPath, BaseModel, Json
  12. from pydantic._internal._typing_extra import origin_is_union
  13. from pydantic._internal._utils import deep_update, lenient_issubclass
  14. from pydantic.fields import FieldInfo
  15. from typing_extensions import get_args, get_origin
  16. from pydantic_settings.utils import path_type_label
  17. if TYPE_CHECKING:
  18. from pydantic_settings.main import BaseSettings
  19. DotenvType = Union[Path, str, List[Union[Path, str]], Tuple[Union[Path, str], ...]]
  20. # This is used as default value for `_env_file` in the `BaseSettings` class and
  21. # `env_file` in `DotEnvSettingsSource` so the default can be distinguished from `None`.
  22. # See the docstring of `BaseSettings` for more details.
  23. ENV_FILE_SENTINEL: DotenvType = Path('')
  24. class SettingsError(ValueError):
  25. pass
  26. class PydanticBaseSettingsSource(ABC):
  27. """
  28. Abstract base class for settings sources, every settings source classes should inherit from it.
  29. """
  30. def __init__(self, settings_cls: type[BaseSettings]):
  31. self.settings_cls = settings_cls
  32. self.config = settings_cls.model_config
  33. @abstractmethod
  34. def get_field_value(self, field: FieldInfo, field_name: str) -> tuple[Any, str, bool]:
  35. """
  36. Gets the value, the key for model creation, and a flag to determine whether value is complex.
  37. This is an abstract method that should be overridden in every settings source classes.
  38. Args:
  39. field: The field.
  40. field_name: The field name.
  41. Returns:
  42. A tuple contains the key, value and a flag to determine whether value is complex.
  43. """
  44. pass
  45. def field_is_complex(self, field: FieldInfo) -> bool:
  46. """
  47. Checks whether a field is complex, in which case it will attempt to be parsed as JSON.
  48. Args:
  49. field: The field.
  50. Returns:
  51. Whether the field is complex.
  52. """
  53. return _annotation_is_complex(field.annotation, field.metadata)
  54. def prepare_field_value(self, field_name: str, field: FieldInfo, value: Any, value_is_complex: bool) -> Any:
  55. """
  56. Prepares the value of a field.
  57. Args:
  58. field_name: The field name.
  59. field: The field.
  60. value: The value of the field that has to be prepared.
  61. value_is_complex: A flag to determine whether value is complex.
  62. Returns:
  63. The prepared value.
  64. """
  65. if value is not None and (self.field_is_complex(field) or value_is_complex):
  66. return self.decode_complex_value(field_name, field, value)
  67. return value
  68. def decode_complex_value(self, field_name: str, field: FieldInfo, value: Any) -> Any:
  69. """
  70. Decode the value for a complex field
  71. Args:
  72. field_name: The field name.
  73. field: The field.
  74. value: The value of the field that has to be prepared.
  75. Returns:
  76. The decoded value for further preparation
  77. """
  78. return json.loads(value)
  79. @abstractmethod
  80. def __call__(self) -> dict[str, Any]:
  81. pass
  82. class InitSettingsSource(PydanticBaseSettingsSource):
  83. """
  84. Source class for loading values provided during settings class initialization.
  85. """
  86. def __init__(self, settings_cls: type[BaseSettings], init_kwargs: dict[str, Any]):
  87. self.init_kwargs = init_kwargs
  88. super().__init__(settings_cls)
  89. def get_field_value(self, field: FieldInfo, field_name: str) -> tuple[Any, str, bool]:
  90. # Nothing to do here. Only implement the return statement to make mypy happy
  91. return None, '', False
  92. def __call__(self) -> dict[str, Any]:
  93. return self.init_kwargs
  94. def __repr__(self) -> str:
  95. return f'InitSettingsSource(init_kwargs={self.init_kwargs!r})'
  96. class PydanticBaseEnvSettingsSource(PydanticBaseSettingsSource):
  97. def __init__(
  98. self, settings_cls: type[BaseSettings], case_sensitive: bool | None = None, env_prefix: str | None = None
  99. ) -> None:
  100. super().__init__(settings_cls)
  101. self.case_sensitive = case_sensitive if case_sensitive is not None else self.config.get('case_sensitive', False)
  102. self.env_prefix = env_prefix if env_prefix is not None else self.config.get('env_prefix', '')
  103. def _apply_case_sensitive(self, value: str) -> str:
  104. return value.lower() if not self.case_sensitive else value
  105. def _extract_field_info(self, field: FieldInfo, field_name: str) -> list[tuple[str, str, bool]]:
  106. """
  107. Extracts field info. This info is used to get the value of field from environment variables.
  108. It returns a list of tuples, each tuple contains:
  109. * field_key: The key of field that has to be used in model creation.
  110. * env_name: The environment variable name of the field.
  111. * value_is_complex: A flag to determine whether the value from environment variable
  112. is complex and has to be parsed.
  113. Args:
  114. field (FieldInfo): The field.
  115. field_name (str): The field name.
  116. Returns:
  117. list[tuple[str, str, bool]]: List of tuples, each tuple contains field_key, env_name, and value_is_complex.
  118. """
  119. field_info: list[tuple[str, str, bool]] = []
  120. if isinstance(field.validation_alias, (AliasChoices, AliasPath)):
  121. v_alias: str | list[str | int] | list[list[str | int]] | None = field.validation_alias.convert_to_aliases()
  122. else:
  123. v_alias = field.validation_alias
  124. if v_alias:
  125. if isinstance(v_alias, list): # AliasChoices, AliasPath
  126. for alias in v_alias:
  127. if isinstance(alias, str): # AliasPath
  128. field_info.append((alias, self._apply_case_sensitive(alias), True if len(alias) > 1 else False))
  129. elif isinstance(alias, list): # AliasChoices
  130. first_arg = cast(str, alias[0]) # first item of an AliasChoices must be a str
  131. field_info.append(
  132. (first_arg, self._apply_case_sensitive(first_arg), True if len(alias) > 1 else False)
  133. )
  134. else: # string validation alias
  135. field_info.append((v_alias, self._apply_case_sensitive(v_alias), False))
  136. else:
  137. field_info.append((field_name, self._apply_case_sensitive(self.env_prefix + field_name), False))
  138. return field_info
  139. def _replace_field_names_case_insensitively(self, field: FieldInfo, field_values: dict[str, Any]) -> dict[str, Any]:
  140. """
  141. Replace field names in values dict by looking in models fields insensitively.
  142. By having the following models:
  143. ```py
  144. class SubSubSub(BaseModel):
  145. VaL3: str
  146. class SubSub(BaseModel):
  147. Val2: str
  148. SUB_sub_SuB: SubSubSub
  149. class Sub(BaseModel):
  150. VAL1: str
  151. SUB_sub: SubSub
  152. class Settings(BaseSettings):
  153. nested: Sub
  154. model_config = SettingsConfigDict(env_nested_delimiter='__')
  155. ```
  156. Then:
  157. _replace_field_names_case_insensitively(
  158. field,
  159. {"val1": "v1", "sub_SUB": {"VAL2": "v2", "sub_SUB_sUb": {"vAl3": "v3"}}}
  160. )
  161. Returns {'VAL1': 'v1', 'SUB_sub': {'Val2': 'v2', 'SUB_sub_SuB': {'VaL3': 'v3'}}}
  162. """
  163. values: dict[str, Any] = {}
  164. for name, value in field_values.items():
  165. sub_model_field: FieldInfo | None = None
  166. # This is here to make mypy happy
  167. # Item "None" of "Optional[Type[Any]]" has no attribute "model_fields"
  168. if not field.annotation or not hasattr(field.annotation, 'model_fields'):
  169. values[name] = value
  170. continue
  171. # Find field in sub model by looking in fields case insensitively
  172. for sub_model_field_name, f in field.annotation.model_fields.items():
  173. if not f.validation_alias and sub_model_field_name.lower() == name.lower():
  174. sub_model_field = f
  175. break
  176. if not sub_model_field:
  177. values[name] = value
  178. continue
  179. if lenient_issubclass(sub_model_field.annotation, BaseModel) and isinstance(value, dict):
  180. values[sub_model_field_name] = self._replace_field_names_case_insensitively(sub_model_field, value)
  181. else:
  182. values[sub_model_field_name] = value
  183. return values
  184. def __call__(self) -> dict[str, Any]:
  185. data: dict[str, Any] = {}
  186. for field_name, field in self.settings_cls.model_fields.items():
  187. try:
  188. field_value, field_key, value_is_complex = self.get_field_value(field, field_name)
  189. except Exception as e:
  190. raise SettingsError(
  191. f'error getting value for field "{field_name}" from source "{self.__class__.__name__}"'
  192. ) from e
  193. try:
  194. field_value = self.prepare_field_value(field_name, field, field_value, value_is_complex)
  195. except ValueError as e:
  196. raise SettingsError(
  197. f'error parsing value for field "{field_name}" from source "{self.__class__.__name__}"'
  198. ) from e
  199. if field_value is not None:
  200. if (
  201. not self.case_sensitive
  202. and lenient_issubclass(field.annotation, BaseModel)
  203. and isinstance(field_value, dict)
  204. ):
  205. data[field_key] = self._replace_field_names_case_insensitively(field, field_value)
  206. else:
  207. data[field_key] = field_value
  208. return data
  209. class SecretsSettingsSource(PydanticBaseEnvSettingsSource):
  210. """
  211. Source class for loading settings values from secret files.
  212. """
  213. def __init__(
  214. self,
  215. settings_cls: type[BaseSettings],
  216. secrets_dir: str | Path | None = None,
  217. case_sensitive: bool | None = None,
  218. env_prefix: str | None = None,
  219. ) -> None:
  220. super().__init__(settings_cls, case_sensitive, env_prefix)
  221. self.secrets_dir = secrets_dir if secrets_dir is not None else self.config.get('secrets_dir')
  222. def __call__(self) -> dict[str, Any]:
  223. """
  224. Build fields from "secrets" files.
  225. """
  226. secrets: dict[str, str | None] = {}
  227. if self.secrets_dir is None:
  228. return secrets
  229. self.secrets_path = Path(self.secrets_dir).expanduser()
  230. if not self.secrets_path.exists():
  231. warnings.warn(f'directory "{self.secrets_path}" does not exist')
  232. return secrets
  233. if not self.secrets_path.is_dir():
  234. raise SettingsError(f'secrets_dir must reference a directory, not a {path_type_label(self.secrets_path)}')
  235. return super().__call__()
  236. @classmethod
  237. def find_case_path(cls, dir_path: Path, file_name: str, case_sensitive: bool) -> Path | None:
  238. """
  239. Find a file within path's directory matching filename, optionally ignoring case.
  240. Args:
  241. dir_path: Directory path.
  242. file_name: File name.
  243. case_sensitive: Whether to search for file name case sensitively.
  244. Returns:
  245. Whether file path or `None` if file does not exist in directory.
  246. """
  247. for f in dir_path.iterdir():
  248. if f.name == file_name:
  249. return f
  250. elif not case_sensitive and f.name.lower() == file_name.lower():
  251. return f
  252. return None
  253. def get_field_value(self, field: FieldInfo, field_name: str) -> tuple[Any, str, bool]:
  254. """
  255. Gets the value for field from secret file and a flag to determine whether value is complex.
  256. Args:
  257. field: The field.
  258. field_name: The field name.
  259. Returns:
  260. A tuple contains the key, value if the file exists otherwise `None`, and
  261. a flag to determine whether value is complex.
  262. """
  263. for field_key, env_name, value_is_complex in self._extract_field_info(field, field_name):
  264. path = self.find_case_path(self.secrets_path, env_name, self.case_sensitive)
  265. if not path:
  266. # path does not exist, we currently don't return a warning for this
  267. continue
  268. if path.is_file():
  269. return path.read_text().strip(), field_key, value_is_complex
  270. else:
  271. warnings.warn(
  272. f'attempted to load secret file "{path}" but found a {path_type_label(path)} instead.',
  273. stacklevel=4,
  274. )
  275. return None, field_key, value_is_complex
  276. def __repr__(self) -> str:
  277. return f'SecretsSettingsSource(secrets_dir={self.secrets_dir!r})'
  278. class EnvSettingsSource(PydanticBaseEnvSettingsSource):
  279. """
  280. Source class for loading settings values from environment variables.
  281. """
  282. def __init__(
  283. self,
  284. settings_cls: type[BaseSettings],
  285. case_sensitive: bool | None = None,
  286. env_prefix: str | None = None,
  287. env_nested_delimiter: str | None = None,
  288. ) -> None:
  289. super().__init__(settings_cls, case_sensitive, env_prefix)
  290. self.env_nested_delimiter = (
  291. env_nested_delimiter if env_nested_delimiter is not None else self.config.get('env_nested_delimiter')
  292. )
  293. self.env_prefix_len = len(self.env_prefix)
  294. self.env_vars = self._load_env_vars()
  295. def _load_env_vars(self) -> Mapping[str, str | None]:
  296. if self.case_sensitive:
  297. return os.environ
  298. return {k.lower(): v for k, v in os.environ.items()}
  299. def get_field_value(self, field: FieldInfo, field_name: str) -> tuple[Any, str, bool]:
  300. """
  301. Gets the value for field from environment variables and a flag to determine whether value is complex.
  302. Args:
  303. field: The field.
  304. field_name: The field name.
  305. Returns:
  306. A tuple contains the key, value if the file exists otherwise `None`, and
  307. a flag to determine whether value is complex.
  308. """
  309. env_val: str | None = None
  310. for field_key, env_name, value_is_complex in self._extract_field_info(field, field_name):
  311. env_val = self.env_vars.get(env_name)
  312. if env_val is not None:
  313. break
  314. return env_val, field_key, value_is_complex
  315. def prepare_field_value(self, field_name: str, field: FieldInfo, value: Any, value_is_complex: bool) -> Any:
  316. """
  317. Prepare value for the field.
  318. * Extract value for nested field.
  319. * Deserialize value to python object for complex field.
  320. Args:
  321. field: The field.
  322. field_name: The field name.
  323. Returns:
  324. A tuple contains prepared value for the field.
  325. Raises:
  326. ValuesError: When There is an error in deserializing value for complex field.
  327. """
  328. is_complex, allow_parse_failure = self._field_is_complex(field)
  329. if is_complex or value_is_complex:
  330. if value is None:
  331. # field is complex but no value found so far, try explode_env_vars
  332. env_val_built = self.explode_env_vars(field_name, field, self.env_vars)
  333. if env_val_built:
  334. return env_val_built
  335. else:
  336. # field is complex and there's a value, decode that as JSON, then add explode_env_vars
  337. try:
  338. value = self.decode_complex_value(field_name, field, value)
  339. except ValueError as e:
  340. if not allow_parse_failure:
  341. raise e
  342. if isinstance(value, dict):
  343. return deep_update(value, self.explode_env_vars(field_name, field, self.env_vars))
  344. else:
  345. return value
  346. elif value is not None:
  347. # simplest case, field is not complex, we only need to add the value if it was found
  348. return value
  349. def _union_is_complex(self, annotation: type[Any] | None, metadata: list[Any]) -> bool:
  350. return any(_annotation_is_complex(arg, metadata) for arg in get_args(annotation))
  351. def _field_is_complex(self, field: FieldInfo) -> tuple[bool, bool]:
  352. """
  353. Find out if a field is complex, and if so whether JSON errors should be ignored
  354. """
  355. if self.field_is_complex(field):
  356. allow_parse_failure = False
  357. elif origin_is_union(get_origin(field.annotation)) and self._union_is_complex(field.annotation, field.metadata):
  358. allow_parse_failure = True
  359. else:
  360. return False, False
  361. return True, allow_parse_failure
  362. @staticmethod
  363. def next_field(field: FieldInfo | None, key: str) -> FieldInfo | None:
  364. """
  365. Find the field in a sub model by key(env name)
  366. By having the following models:
  367. ```py
  368. class SubSubModel(BaseSettings):
  369. dvals: Dict
  370. class SubModel(BaseSettings):
  371. vals: list[str]
  372. sub_sub_model: SubSubModel
  373. class Cfg(BaseSettings):
  374. sub_model: SubModel
  375. ```
  376. Then:
  377. next_field(sub_model, 'vals') Returns the `vals` field of `SubModel` class
  378. next_field(sub_model, 'sub_sub_model') Returns `sub_sub_model` field of `SubModel` class
  379. Args:
  380. field: The field.
  381. key: The key (env name).
  382. Returns:
  383. Field if it finds the next field otherwise `None`.
  384. """
  385. if not field or origin_is_union(get_origin(field.annotation)):
  386. # no support for Unions of complex BaseSettings fields
  387. return None
  388. elif field.annotation and hasattr(field.annotation, 'model_fields') and field.annotation.model_fields.get(key):
  389. return field.annotation.model_fields[key]
  390. return None
  391. def explode_env_vars(self, field_name: str, field: FieldInfo, env_vars: Mapping[str, str | None]) -> dict[str, Any]:
  392. """
  393. Process env_vars and extract the values of keys containing env_nested_delimiter into nested dictionaries.
  394. This is applied to a single field, hence filtering by env_var prefix.
  395. Args:
  396. field_name: The field name.
  397. field: The field.
  398. env_vars: Environment variables.
  399. Returns:
  400. A dictionaty contains extracted values from nested env values.
  401. """
  402. prefixes = [
  403. f'{env_name}{self.env_nested_delimiter}' for _, env_name, _ in self._extract_field_info(field, field_name)
  404. ]
  405. result: dict[str, Any] = {}
  406. for env_name, env_val in env_vars.items():
  407. if not any(env_name.startswith(prefix) for prefix in prefixes):
  408. continue
  409. # we remove the prefix before splitting in case the prefix has characters in common with the delimiter
  410. env_name_without_prefix = env_name[self.env_prefix_len :]
  411. _, *keys, last_key = env_name_without_prefix.split(self.env_nested_delimiter)
  412. env_var = result
  413. target_field: FieldInfo | None = field
  414. for key in keys:
  415. target_field = self.next_field(target_field, key)
  416. env_var = env_var.setdefault(key, {})
  417. # get proper field with last_key
  418. target_field = self.next_field(target_field, last_key)
  419. # check if env_val maps to a complex field and if so, parse the env_val
  420. if target_field and env_val:
  421. is_complex, allow_json_failure = self._field_is_complex(target_field)
  422. if is_complex:
  423. try:
  424. env_val = self.decode_complex_value(last_key, target_field, env_val)
  425. except ValueError as e:
  426. if not allow_json_failure:
  427. raise e
  428. env_var[last_key] = env_val
  429. return result
  430. def __repr__(self) -> str:
  431. return (
  432. f'EnvSettingsSource(env_nested_delimiter={self.env_nested_delimiter!r}, '
  433. f'env_prefix_len={self.env_prefix_len!r})'
  434. )
  435. class DotEnvSettingsSource(EnvSettingsSource):
  436. """
  437. Source class for loading settings values from env files.
  438. """
  439. def __init__(
  440. self,
  441. settings_cls: type[BaseSettings],
  442. env_file: DotenvType | None = ENV_FILE_SENTINEL,
  443. env_file_encoding: str | None = None,
  444. case_sensitive: bool | None = None,
  445. env_prefix: str | None = None,
  446. env_nested_delimiter: str | None = None,
  447. ) -> None:
  448. self.env_file = env_file if env_file != ENV_FILE_SENTINEL else settings_cls.model_config.get('env_file')
  449. self.env_file_encoding = (
  450. env_file_encoding if env_file_encoding is not None else settings_cls.model_config.get('env_file_encoding')
  451. )
  452. super().__init__(settings_cls, case_sensitive, env_prefix, env_nested_delimiter)
  453. def _load_env_vars(self) -> Mapping[str, str | None]:
  454. return self._read_env_files(self.case_sensitive)
  455. def _read_env_files(self, case_sensitive: bool) -> Mapping[str, str | None]:
  456. env_files = self.env_file
  457. if env_files is None:
  458. return {}
  459. if isinstance(env_files, (str, os.PathLike)):
  460. env_files = [env_files]
  461. dotenv_vars: dict[str, str | None] = {}
  462. for env_file in env_files:
  463. env_path = Path(env_file).expanduser()
  464. if env_path.is_file():
  465. dotenv_vars.update(
  466. read_env_file(env_path, encoding=self.env_file_encoding, case_sensitive=case_sensitive)
  467. )
  468. return dotenv_vars
  469. def __call__(self) -> dict[str, Any]:
  470. data: dict[str, Any] = super().__call__()
  471. data_lower_keys: list[str] = []
  472. if not self.case_sensitive:
  473. data_lower_keys = [x.lower() for x in data.keys()]
  474. # As `extra` config is allowed in dotenv settings source, We have to
  475. # update data with extra env variabels from dotenv file.
  476. for env_name, env_value in self.env_vars.items():
  477. if env_name.startswith(self.env_prefix) and env_value is not None:
  478. env_name_without_prefix = env_name[self.env_prefix_len :]
  479. first_key, *_ = env_name_without_prefix.split(self.env_nested_delimiter)
  480. if (data_lower_keys and first_key not in data_lower_keys) or (
  481. not data_lower_keys and first_key not in data
  482. ):
  483. data[first_key] = env_value
  484. return data
  485. def __repr__(self) -> str:
  486. return (
  487. f'DotEnvSettingsSource(env_file={self.env_file!r}, env_file_encoding={self.env_file_encoding!r}, '
  488. f'env_nested_delimiter={self.env_nested_delimiter!r}, env_prefix_len={self.env_prefix_len!r})'
  489. )
  490. def read_env_file(
  491. file_path: Path, *, encoding: str | None = None, case_sensitive: bool = False
  492. ) -> Mapping[str, str | None]:
  493. file_vars: dict[str, str | None] = dotenv_values(file_path, encoding=encoding or 'utf8')
  494. if not case_sensitive:
  495. return {k.lower(): v for k, v in file_vars.items()}
  496. else:
  497. return file_vars
  498. def _annotation_is_complex(annotation: type[Any] | None, metadata: list[Any]) -> bool:
  499. if any(isinstance(md, Json) for md in metadata): # type: ignore[misc]
  500. return False
  501. origin = get_origin(annotation)
  502. return (
  503. _annotation_is_complex_inner(annotation)
  504. or _annotation_is_complex_inner(origin)
  505. or hasattr(origin, '__pydantic_core_schema__')
  506. or hasattr(origin, '__get_pydantic_core_schema__')
  507. )
  508. def _annotation_is_complex_inner(annotation: type[Any] | None) -> bool:
  509. if lenient_issubclass(annotation, (str, bytes)):
  510. return False
  511. return lenient_issubclass(annotation, (BaseModel, Mapping, Sequence, tuple, set, frozenset, deque)) or is_dataclass(
  512. annotation
  513. )