100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
"""TOML file settings source."""
|
|
|
|
from __future__ import annotations as _annotations
|
|
|
|
import sys
|
|
from collections.abc import Sequence
|
|
from pathlib import Path
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
)
|
|
|
|
from ..base import ConfigFileSourceMixin, InitSettingsSource
|
|
from ..types import DEFAULT_PATH, ConfigFileSourceType
|
|
from ..utils import InitState
|
|
|
|
if TYPE_CHECKING:
|
|
from pydantic_settings.main import BaseSettings
|
|
|
|
from ..types import Traversable
|
|
|
|
if sys.version_info >= (3, 11):
|
|
import tomllib
|
|
else:
|
|
tomllib = None
|
|
import tomli
|
|
else:
|
|
tomllib = None
|
|
tomli = None
|
|
|
|
|
|
def import_toml() -> None:
|
|
global tomli
|
|
global tomllib
|
|
if sys.version_info < (3, 11):
|
|
if tomli is not None:
|
|
return
|
|
try:
|
|
import tomli
|
|
except ImportError as e: # pragma: no cover
|
|
raise ImportError('tomli is not installed, run `pip install pydantic-settings[toml]`') from e
|
|
else:
|
|
if tomllib is not None:
|
|
return
|
|
import tomllib
|
|
|
|
|
|
class TomlConfigSettingsSource(InitSettingsSource, ConfigFileSourceMixin):
|
|
"""
|
|
A source class that loads variables from a TOML file
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
settings_cls: type[BaseSettings],
|
|
toml_file: ConfigFileSourceType | None = DEFAULT_PATH,
|
|
toml_table_header: tuple[str, ...] = (),
|
|
deep_merge: bool = False,
|
|
_init_state: InitState | None = None,
|
|
):
|
|
self.toml_file_path = toml_file if toml_file != DEFAULT_PATH else settings_cls.model_config.get('toml_file')
|
|
self.toml_table_header = (
|
|
toml_table_header if toml_table_header else settings_cls.model_config.get('toml_table_header', ())
|
|
)
|
|
self.toml_data = self._read_files(self.toml_file_path, deep_merge=deep_merge)
|
|
|
|
if self._any_file_exists(self.toml_file_path):
|
|
for key in self.toml_table_header:
|
|
if key not in self.toml_data:
|
|
raise KeyError(f'toml_table_header key "{key}" not found in {self.toml_file_path}')
|
|
self.toml_data = self.toml_data[key]
|
|
|
|
super().__init__(settings_cls, self.toml_data, _init_state=_init_state)
|
|
|
|
def _read_file(self, file_path: Path | Traversable) -> dict[str, Any]:
|
|
import_toml()
|
|
with file_path.open(mode='rb') as toml_file:
|
|
if sys.version_info < (3, 11):
|
|
return tomli.load(toml_file)
|
|
return tomllib.load(toml_file)
|
|
|
|
@staticmethod
|
|
def _any_file_exists(paths: ConfigFileSourceType | None) -> bool:
|
|
"""Check if any of the given file paths exist."""
|
|
if paths is None:
|
|
return False
|
|
if isinstance(paths, str) or not isinstance(paths, Sequence):
|
|
paths = [paths]
|
|
|
|
def _exists(path: Path | str | Traversable) -> bool:
|
|
if isinstance(path, (str, Path)):
|
|
return Path(path).exists()
|
|
# Non-`Path` `Traversable` (e.g. a resource inside a zip/wheel) is not
|
|
# `os.PathLike`, so it can't be wrapped in `Path`; query it directly.
|
|
return path.is_file()
|
|
|
|
return any(_exists(path) for path in paths)
|
|
|
|
def __repr__(self) -> str:
|
|
return f'{self.__class__.__name__}(toml_file={self.toml_file_path}, toml_table_header={self.toml_table_header})'
|