272 lines
9.6 KiB
Python
272 lines
9.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections.abc import Mapping
|
||
|
|
from enum import Enum
|
||
|
|
from pathlib import Path
|
||
|
|
from types import UnionType
|
||
|
|
from typing import Annotated, Any, Literal, TypeVar, Union, cast, get_args, get_origin
|
||
|
|
|
||
|
|
from pydantic import BaseModel, ValidationError
|
||
|
|
import tomli_w
|
||
|
|
|
||
|
|
from core_lib.common.exception import ModelRegistrationError, ConfigValidationError, SettingsError
|
||
|
|
|
||
|
|
try:
|
||
|
|
import tomllib
|
||
|
|
except ModuleNotFoundError: # Python 3.10
|
||
|
|
import tomli as tomllib
|
||
|
|
|
||
|
|
ModelT = TypeVar("ModelT", bound=BaseModel)
|
||
|
|
|
||
|
|
|
||
|
|
class SettingsManager:
|
||
|
|
"""
|
||
|
|
A settings manager (Based-on pydantic)
|
||
|
|
|
||
|
|
Vibed!
|
||
|
|
"""
|
||
|
|
def __init__(self, path: str | Path | None = None) -> None:
|
||
|
|
self.path = Path(path) if path is not None else None
|
||
|
|
self._models: dict[str, type[BaseModel]] = {}
|
||
|
|
self._settings: dict[str, BaseModel] = {}
|
||
|
|
|
||
|
|
@property
|
||
|
|
def registered_sections(self) -> tuple[str, ...]:
|
||
|
|
return tuple(self._models)
|
||
|
|
|
||
|
|
def register(self, section: str, model: type[ModelT], *, replace: bool = False) -> type[ModelT]:
|
||
|
|
"""
|
||
|
|
Register a new model
|
||
|
|
:param section:
|
||
|
|
:param model:
|
||
|
|
:param replace:
|
||
|
|
:return:
|
||
|
|
model: type[ModelT]
|
||
|
|
"""
|
||
|
|
if not section:
|
||
|
|
raise ModelRegistrationError(
|
||
|
|
"Section must be a non-empty"
|
||
|
|
)
|
||
|
|
elif "." in section:
|
||
|
|
raise ModelRegistrationError(
|
||
|
|
"Top level section cannot contain dots"
|
||
|
|
)
|
||
|
|
|
||
|
|
if not isinstance(model, type) or not issubclass(model, BaseModel):
|
||
|
|
raise ModelRegistrationError("Model must be a Pydantic BaseModel class")
|
||
|
|
|
||
|
|
if section in self._models and not replace:
|
||
|
|
raise ModelRegistrationError(f"Section {section!r} is already registered")
|
||
|
|
|
||
|
|
self._models[section] = model
|
||
|
|
self._settings.pop(section, None)
|
||
|
|
return model
|
||
|
|
|
||
|
|
def model(self, section: str):
|
||
|
|
def decorator(model: type[ModelT]) -> type[ModelT]:
|
||
|
|
return self.register(section, model)
|
||
|
|
|
||
|
|
return decorator
|
||
|
|
|
||
|
|
def load(self, path: str | Path | None = None, *, reject_unknown: bool = True) -> SettingsManager:
|
||
|
|
"""
|
||
|
|
Load a toml settings file
|
||
|
|
:param reject_unknown: Raise exception if settings file contains unknown section
|
||
|
|
:param path: If not set, use default path instead
|
||
|
|
"""
|
||
|
|
config_path = self._resolve_path(path)
|
||
|
|
|
||
|
|
# Check file exist
|
||
|
|
if not config_path.exists() or not config_path.is_file():
|
||
|
|
raise FileNotFoundError(f"Settings file {config_path} does not exist")
|
||
|
|
|
||
|
|
with config_path.open("rb") as file:
|
||
|
|
raw = tomllib.load(file)
|
||
|
|
|
||
|
|
unknown = set(raw).difference(self._models)
|
||
|
|
if reject_unknown and unknown:
|
||
|
|
names = ", ".join(sorted(unknown))
|
||
|
|
raise SettingsError(f"unregistered TOML section(s): {names}")
|
||
|
|
|
||
|
|
loaded: dict[str, BaseModel] = {}
|
||
|
|
errors: dict[str, ValidationError] = {}
|
||
|
|
for section, model in self._models.items():
|
||
|
|
value: Any = raw.get(section, {})
|
||
|
|
try:
|
||
|
|
# Validate and save section data
|
||
|
|
loaded[section] = model.model_validate(value)
|
||
|
|
except ValidationError as error:
|
||
|
|
errors[section] = error
|
||
|
|
|
||
|
|
if errors:
|
||
|
|
raise ConfigValidationError(errors)
|
||
|
|
|
||
|
|
self.path = config_path
|
||
|
|
self._settings = loaded
|
||
|
|
return self
|
||
|
|
|
||
|
|
def get(self, section: str, model: type[ModelT] | None = None) -> ModelT:
|
||
|
|
"""
|
||
|
|
Get target section model
|
||
|
|
:param section:
|
||
|
|
:param model:
|
||
|
|
:return:
|
||
|
|
model: type[ModelT]
|
||
|
|
"""
|
||
|
|
if section not in self._models:
|
||
|
|
raise KeyError(f"settings section {section!r} is not registered")
|
||
|
|
if section not in self._settings:
|
||
|
|
raise SettingsError(f"settings have not been loaded for {section!r}")
|
||
|
|
|
||
|
|
value = self._settings[section]
|
||
|
|
if model is not None and not isinstance(value, model):
|
||
|
|
raise TypeError(
|
||
|
|
f"section {section!r} contains {type(value).__name__}, "
|
||
|
|
f"not {model.__name__}"
|
||
|
|
)
|
||
|
|
|
||
|
|
return cast(ModelT, value)
|
||
|
|
|
||
|
|
def set(self, section: str, value: BaseModel | Mapping[str, Any]) -> BaseModel:
|
||
|
|
"""Validate and replace one loaded settings section."""
|
||
|
|
try:
|
||
|
|
model = self._models[section]
|
||
|
|
except KeyError as error:
|
||
|
|
raise KeyError(f"settings section {section!r} is not registered") from error
|
||
|
|
|
||
|
|
validated = model.model_validate(value)
|
||
|
|
self._settings[section] = validated
|
||
|
|
return validated
|
||
|
|
|
||
|
|
def load_defaults(self) -> SettingsManager:
|
||
|
|
"""Load every registered model using its declared default values."""
|
||
|
|
loaded: dict[str, BaseModel] = {}
|
||
|
|
errors: dict[str, ValidationError] = {}
|
||
|
|
for section, model in self._models.items():
|
||
|
|
try:
|
||
|
|
loaded[section] = model.model_validate({})
|
||
|
|
except ValidationError as error:
|
||
|
|
errors[section] = error
|
||
|
|
if errors:
|
||
|
|
raise ConfigValidationError(errors)
|
||
|
|
self._settings = loaded
|
||
|
|
return self
|
||
|
|
|
||
|
|
def save(self, path: str | Path | None = None) -> Path:
|
||
|
|
"""Validate all loaded sections and save them as UTF-8 TOML."""
|
||
|
|
config_path = self._resolve_path(path)
|
||
|
|
missing = set(self._models).difference(self._settings)
|
||
|
|
if missing:
|
||
|
|
names = ", ".join(sorted(missing))
|
||
|
|
raise SettingsError(f"settings section(s) have not been loaded: {names}")
|
||
|
|
|
||
|
|
data = {
|
||
|
|
section: self._models[section]
|
||
|
|
.model_validate(self._settings[section])
|
||
|
|
.model_dump(mode="json")
|
||
|
|
for section in self._models
|
||
|
|
}
|
||
|
|
config_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
temporary_path = config_path.with_suffix(config_path.suffix + ".tmp")
|
||
|
|
temporary_path.write_bytes(tomli_w.dumps(data).encode("utf-8"))
|
||
|
|
temporary_path.replace(config_path)
|
||
|
|
self.path = config_path
|
||
|
|
return config_path
|
||
|
|
|
||
|
|
def choices(self, section: str, field: str) -> tuple[Any, ...]:
|
||
|
|
"""
|
||
|
|
Return choices for a field, accepting paths such as ``window.mode``.
|
||
|
|
|
||
|
|
This can be used by a settings page to populate a combo box. Pydantic
|
||
|
|
uses the same annotation to validate values read from TOML.
|
||
|
|
:param section:
|
||
|
|
:param field:
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
model = self._models[section]
|
||
|
|
except KeyError as error:
|
||
|
|
raise KeyError(f"settings section {section!r} is not registered") from error
|
||
|
|
|
||
|
|
annotation: Any = model
|
||
|
|
for part in field.split("."):
|
||
|
|
nested_model = _model_type(annotation)
|
||
|
|
if nested_model is None:
|
||
|
|
raise KeyError(
|
||
|
|
f"{part!r} cannot be resolved in field path {section}.{field}"
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
annotation = nested_model.model_fields[part].annotation
|
||
|
|
except KeyError as error:
|
||
|
|
raise KeyError(f"field path {section}.{field} does not exist") from error
|
||
|
|
|
||
|
|
values = _choice_values(annotation)
|
||
|
|
if values is None:
|
||
|
|
raise SettingsError(
|
||
|
|
f"field {section}.{field} is not declared with Literal or Enum"
|
||
|
|
)
|
||
|
|
return values
|
||
|
|
|
||
|
|
def validate(self, data: Mapping[str, Any], *, reject_unknown: bool = True) -> None:
|
||
|
|
"""
|
||
|
|
Validate an in-memory mapping without changing loaded settings.
|
||
|
|
:param data:
|
||
|
|
:param reject_unknown:
|
||
|
|
:return:
|
||
|
|
"""
|
||
|
|
unknown = set(data).difference(self._models)
|
||
|
|
if reject_unknown and unknown:
|
||
|
|
names = ", ".join(sorted(unknown))
|
||
|
|
raise SettingsError(f"unregistered settings section(s): {names}")
|
||
|
|
|
||
|
|
errors: dict[str, ValidationError] = {}
|
||
|
|
for section, model in self._models.items():
|
||
|
|
try:
|
||
|
|
model.model_validate(data.get(section, {}))
|
||
|
|
except ValidationError as error:
|
||
|
|
errors[section] = error
|
||
|
|
if errors:
|
||
|
|
raise ConfigValidationError(errors)
|
||
|
|
|
||
|
|
def _resolve_path(self, path: str | Path | None) -> Path:
|
||
|
|
resolved = Path(path) if path is not None else self.path
|
||
|
|
|
||
|
|
if resolved is None:
|
||
|
|
raise SettingsError("no TOML configuration path was provided")
|
||
|
|
|
||
|
|
return resolved
|
||
|
|
|
||
|
|
def _choice_values(annotation: Any) -> tuple[Any, ...] | None:
|
||
|
|
"""Extract choices from a field annotation, including Annotated/Optional."""
|
||
|
|
origin = get_origin(annotation)
|
||
|
|
if origin is Literal:
|
||
|
|
return get_args(annotation)
|
||
|
|
if isinstance(annotation, type) and issubclass(annotation, Enum):
|
||
|
|
return tuple(member.value for member in annotation)
|
||
|
|
if origin is Annotated:
|
||
|
|
return _choice_values(get_args(annotation)[0])
|
||
|
|
if origin in (Union, UnionType):
|
||
|
|
choices: list[Any] = []
|
||
|
|
found = False
|
||
|
|
for member in get_args(annotation):
|
||
|
|
values = _choice_values(member)
|
||
|
|
if values is not None:
|
||
|
|
choices.extend(values)
|
||
|
|
found = True
|
||
|
|
return tuple(choices) if found else None
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _model_type(annotation: Any) -> type[BaseModel] | None:
|
||
|
|
"""Find a nested BaseModel type inside Annotated or Optional."""
|
||
|
|
if isinstance(annotation, type) and issubclass(annotation, BaseModel):
|
||
|
|
return annotation
|
||
|
|
origin = get_origin(annotation)
|
||
|
|
if origin is Annotated:
|
||
|
|
return _model_type(get_args(annotation)[0])
|
||
|
|
if origin in (Union, UnionType):
|
||
|
|
for member in get_args(annotation):
|
||
|
|
model = _model_type(member)
|
||
|
|
if model is not None:
|
||
|
|
return model
|
||
|
|
return None
|