f71bfd3b8c
Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections import defaultdict
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from vibe.core.config.layer import (
|
|
ConfigLayer,
|
|
EmptyLayerError,
|
|
RawConfig,
|
|
UntrustedLayerError,
|
|
)
|
|
from vibe.core.config.schema import ConfigFragment, ConfigSchema, MergeFieldMetadata
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _LayerData:
|
|
name: str
|
|
data: dict[str, Any]
|
|
|
|
|
|
class ConfigBuilder[S: ConfigSchema]:
|
|
"""Collects layers and merges them into an immutable Config[S]."""
|
|
|
|
def __init__(self, schema: type[S]) -> None:
|
|
self._schema = schema
|
|
self._layers: list[ConfigLayer[RawConfig]] = []
|
|
self._lock = asyncio.Lock()
|
|
|
|
def add_layer(self, layer: ConfigLayer[RawConfig]) -> None:
|
|
self._layers.append(layer)
|
|
|
|
def add_layers(self, layers: list[ConfigLayer[RawConfig]]) -> None:
|
|
self._layers.extend(layers)
|
|
|
|
@property
|
|
def layers(self) -> list[ConfigLayer[RawConfig]]:
|
|
return self._layers
|
|
|
|
async def build(self, force_load: bool = False) -> S:
|
|
"""Merge all layers and return a validated schema.
|
|
|
|
Untrusted and empty layers are skipped.
|
|
Pass ``force_load=True`` to bypass caching.
|
|
"""
|
|
async with self._lock:
|
|
internal_layers = self._layers.copy()
|
|
|
|
layer_dicts: list[_LayerData] = []
|
|
for layer in internal_layers:
|
|
try:
|
|
data = await layer.load(force=force_load)
|
|
raw = data.model_dump()
|
|
if raw:
|
|
layer_dicts.append(_LayerData(name=layer.name, data=raw))
|
|
except (UntrustedLayerError, EmptyLayerError):
|
|
continue
|
|
|
|
merged, origins = self._merge_fields(self._schema, layer_dicts)
|
|
return self._schema(origins=origins, **merged)
|
|
|
|
def _merge_fields(
|
|
self, schema: type[S], layer_dicts: list[_LayerData]
|
|
) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
accumulated: dict[str, Any] = defaultdict(dict)
|
|
origins: dict[str, Any] = {}
|
|
|
|
for ld in layer_dicts:
|
|
for key, value in ld.data.items():
|
|
if key not in schema.model_fields:
|
|
continue
|
|
|
|
field_info = schema.model_fields[key]
|
|
annotation = field_info.annotation
|
|
if annotation is None:
|
|
continue
|
|
|
|
is_fragment = isinstance(annotation, type) and issubclass(
|
|
annotation, ConfigFragment
|
|
)
|
|
if is_fragment:
|
|
if not isinstance(value, dict):
|
|
continue
|
|
|
|
for fragment_key, fragment_value in value.items():
|
|
if fragment_key not in annotation.model_fields:
|
|
continue
|
|
|
|
fragment_field = annotation.model_fields[fragment_key]
|
|
fragment_meta = MergeFieldMetadata.from_field(fragment_field)
|
|
if fragment_meta is None:
|
|
continue
|
|
|
|
accumulated[key][fragment_key] = (
|
|
fragment_meta.merge_strategy.apply(
|
|
accumulated[key].get(fragment_key),
|
|
fragment_value,
|
|
key_fn=self._make_key_fn(fragment_meta),
|
|
)
|
|
)
|
|
continue
|
|
|
|
meta = MergeFieldMetadata.from_field(field_info)
|
|
if meta is None:
|
|
continue
|
|
|
|
accumulated[key] = meta.merge_strategy.apply(
|
|
accumulated.get(key), value, key_fn=self._make_key_fn(meta)
|
|
)
|
|
|
|
return accumulated, origins
|
|
|
|
def _make_key_fn(
|
|
self, merge_field_meta: MergeFieldMetadata
|
|
) -> Callable[[Any], str] | None:
|
|
merge_key = merge_field_meta.merge_key
|
|
if merge_key is None:
|
|
return None
|
|
|
|
return lambda item: item[merge_key]
|