codekingpro/portable-devtools

1"""TOML file settings source."""2 3from __future__ import annotations as _annotations4 5import sys6from pathlib import Path7from typing import (8    TYPE_CHECKING,9    Any,10)11 12from ..base import ConfigFileSourceMixin, InitSettingsSource13from ..types import DEFAULT_PATH, PathType14 15if TYPE_CHECKING:16    from pydantic_settings.main import BaseSettings17 18    if sys.version_info >= (3, 11):19        import tomllib20    else:21        tomllib = None22    import tomli23else:24    tomllib = None25    tomli = None26 27 28def import_toml() -> None:29    global tomli30    global tomllib31    if sys.version_info < (3, 11):32        if tomli is not None:33            return34        try:35            import tomli36        except ImportError as e:  # pragma: no cover37            raise ImportError('tomli is not installed, run `pip install pydantic-settings[toml]`') from e38    else:39        if tomllib is not None:40            return41        import tomllib42 43 44class TomlConfigSettingsSource(InitSettingsSource, ConfigFileSourceMixin):45    """46    A source class that loads variables from a TOML file47    """48 49    def __init__(50        self,51        settings_cls: type[BaseSettings],52        toml_file: PathType | None = DEFAULT_PATH,53        deep_merge: bool = False,54    ):55        self.toml_file_path = toml_file if toml_file != DEFAULT_PATH else settings_cls.model_config.get('toml_file')56        self.toml_data = self._read_files(self.toml_file_path, deep_merge=deep_merge)57        super().__init__(settings_cls, self.toml_data)58 59    def _read_file(self, file_path: Path) -> dict[str, Any]:60        import_toml()61        with file_path.open(mode='rb') as toml_file:62            if sys.version_info < (3, 11):63                return tomli.load(toml_file)64            return tomllib.load(toml_file)65 66    def __repr__(self) -> str:67        return f'{self.__class__.__name__}(toml_file={self.toml_file_path})'68