Coverage for ml_workbench/config.py: 16%
170 statements
« prev ^ index » next coverage.py v7.11.0, created at 2026-01-06 16:09 +0200
1from __future__ import annotations
3from collections.abc import Iterable, Mapping, MutableMapping
4from copy import deepcopy
5from pathlib import Path
6from string import Formatter
7from typing import Any, Union
9import yaml
11YamlPrimitive = Union[str, int, float, bool, None] # noqa: UP007
12YamlValue = Union[YamlPrimitive, "YamlDict", "YamlList"] # noqa: UP007
13YamlDict = dict[str, YamlValue]
14YamlList = list[YamlValue]
17def _extract_field_names(template: str) -> list[str]:
18 """Return field names used in a ``str.format``-style template string.
20 Example: "{catalog}.{schema}.table" -> ["catalog", "schema"]
21 """
23 field_names: list[str] = []
24 for _literal_text, field_name, _format_spec, _conversion in Formatter().parse(
25 template
26 ):
27 if field_name: # skip None and empty
28 field_names.append(field_name)
29 return field_names
32def _interpolate_string(
33 value: str, variables: Mapping[str, Any], *, strict: bool
34) -> str:
35 """Interpolate placeholders in a string using ``variables``.
37 - If ``strict`` is True (default), raise a KeyError if any placeholder is missing.
38 - If ``strict`` is False, leave unknown placeholders as-is.
39 """
41 if "{" not in value or "}" not in value:
42 return value
44 placeholders = _extract_field_names(value)
45 if not placeholders:
46 return value
48 missing = [name for name in placeholders if name not in variables]
49 if missing and strict:
50 raise KeyError( # noqa: TRY003
51 f"Missing values for placeholders {missing} while formatting: {value!r}"
52 )
54 class _DefaultDict(dict):
55 def __missing__(self, key: str) -> str: # type: ignore[override]
56 return "{" + key + "}"
58 mapping: Mapping[str, Any]
59 mapping = variables if strict else _DefaultDict(variables) # type: ignore[arg-type]
60 return value.format_map(mapping) # type: ignore[arg-type]
63def _interpolate_value(
64 value: YamlValue, variables: Mapping[str, Any], *, strict: bool
65) -> YamlValue:
66 if isinstance(value, str):
67 return _interpolate_string(value, variables, strict=strict)
68 if isinstance(value, dict):
69 return {
70 k: _interpolate_value(v, variables, strict=strict) for k, v in value.items()
71 }
72 if isinstance(value, list):
73 return [_interpolate_value(v, variables, strict=strict) for v in value]
74 return value
77def _wrap(value: YamlValue) -> Any:
78 if isinstance(value, dict):
79 return _ConfigNode(value)
80 if isinstance(value, list):
81 return [_wrap(v) for v in value]
82 return value
85def _unwrap(value: Any) -> YamlValue:
86 if isinstance(value, _ConfigNode):
87 return {k: _unwrap(v) for k, v in value.items()}
88 if isinstance(value, list):
89 return [_unwrap(v) for v in value]
90 return value
93## moved to dataset.py: _is_databricks_table_name, _infer_dataset_format_from_path
96class _ConfigNode(MutableMapping[str, Any]):
97 """Lightweight mapping that supports both dict- and attribute-style access.
99 This is used for nested configuration sections.
100 """
102 def __init__(self, data: YamlDict) -> None:
103 self._data: YamlDict = data
105 # Mapping interface
106 def __getitem__(self, key: str) -> Any:
107 return _wrap(self._data[key])
109 def __setitem__(self, key: str, value: Any) -> None:
110 self._data[key] = _unwrap(value)
112 def __delitem__(self, key: str) -> None:
113 del self._data[key]
115 def __iter__(self) -> Iterable[str]:
116 return iter(self._data)
118 def __len__(self) -> int: # noqa: D401 - trivial
119 return len(self._data)
121 # Attribute-style access
122 def __getattr__(self, name: str) -> Any:
123 try:
124 return _wrap(self._data[name])
125 except KeyError as exc:
126 raise AttributeError(name) from exc # noqa: TRY003
128 def get(self, key: str, default: Any | None = None) -> Any:
129 if key in self._data:
130 return _wrap(self._data[key])
131 return default
133 def to_dict(self) -> YamlDict:
134 return deepcopy(self._data)
136 def __repr__(self) -> str: # pragma: no cover - representation only
137 return f"_ConfigNode({self._data!r})"
140class YamlConfig(_ConfigNode):
141 """Configuration loaded from a YAML file with placeholder interpolation.
143 Placeholders use Python's ``str.format`` style, e.g. "{catalog}".
144 Values for placeholders are assembled as:
145 1) ``defaults`` key from the YAML file (if present), then
146 2) Keyword arguments provided to the constructor, which override defaults.
148 Parameters
149 ----------
150 yaml_path:
151 Path to the YAML file.
152 strict:
153 If True (default), raise when placeholders are missing. If False, leave
154 unknown placeholders unchanged.
155 **variables:
156 Additional key/value pairs to use for placeholder interpolation.
157 """
159 def __init__(
160 self, yaml_path: str | Path, *, strict: bool = True, **variables: Any
161 ) -> None:
162 path = Path(yaml_path)
163 if not path.exists():
164 raise FileNotFoundError(f"YAML file not found: {path}") # noqa: TRY003
166 # Load YAML with recursive include processing
167 loaded = self._load_with_includes(path, set())
169 if not isinstance(loaded, dict):
170 raise TypeError("Top-level YAML structure must be a mapping (dict)") # noqa: TRY003
172 defaults_section: Mapping[str, Any] = (
173 loaded.get("defaults", {})
174 if isinstance(loaded.get("defaults", {}), dict)
175 else {}
176 )
178 # Merge defaults with user-supplied variables (kwargs override defaults)
179 variables_merged: dict[str, Any] = {**defaults_section, **variables}
181 # Interpolate placeholders across the entire loaded YAML structure
182 interpolated: YamlDict = _interpolate_value(
183 loaded, variables_merged, strict=strict
184 ) # type: ignore[assignment]
186 super().__init__(interpolated)
187 self.path: Path = path
188 self.variables: dict[str, Any] = variables_merged
189 self.strict: bool = strict
191 # Post-process datasets to impute missing information
192 # Delayed import to avoid circular dependency
193 from .dataset import Dataset # noqa: PLC0415
195 Dataset.verify_config(self)
197 from .feature import Feature # noqa: PLC0415
199 Feature.verify_config(self)
201 from .model import Model # noqa: PLC0415
203 Model.verify_config(self)
205 from .experiment import Experiment # noqa: PLC0415
207 Experiment.verify_config(self)
209 from .mlflow_conf import MlflowConf # noqa: PLC0415
211 MlflowConf.verify_config(self)
213 def _load_with_includes(
214 self, yaml_path: Path, visited: set[Path]
215 ) -> dict[str, Any]:
216 """
217 Load a YAML file and recursively process any 'include' directives.
219 The 'include' field should be a list of file paths (relative to the current file
220 or absolute). Included files are merged into the current file, with the current
221 file's values taking precedence.
223 Parameters
224 ----------
225 yaml_path : Path
226 Path to the YAML file to load
227 visited : set[Path]
228 Set of already visited paths to prevent circular includes
230 Returns
231 -------
232 Dict[str, Any]
233 Merged dictionary from all included files and the current file
235 Raises
236 ------
237 ValueError
238 If a circular include is detected
239 """
240 # Resolve to absolute path to detect circular includes
241 abs_path = yaml_path.resolve()
243 if abs_path in visited:
244 raise ValueError(f"Circular include detected: {abs_path}") # noqa: TRY003
246 visited.add(abs_path)
248 # Load the current file
249 with abs_path.open("r", encoding="utf-8") as f:
250 current_data: Any = yaml.safe_load(f) or {}
252 if not isinstance(current_data, dict):
253 return current_data
255 # Check for 'include' directive
256 include_list = current_data.pop("include", None)
258 if include_list is None:
259 return current_data
261 # Ensure include is a list
262 if not isinstance(include_list, list):
263 raise TypeError( # noqa: TRY003
264 f"'include' directive must be a list, got {type(include_list).__name__}"
265 )
267 # Start with an empty result
268 result: dict[str, Any] = {}
270 # Process each included file
271 for include_path_str in include_list:
272 if not isinstance(include_path_str, str):
273 raise TypeError( # noqa: TRY003
274 f"Include path must be a string, got {type(include_path_str).__name__}"
275 )
277 # Resolve include path
278 include_path = Path(include_path_str)
279 if not include_path.is_absolute():
280 # Try relative to current file's directory first
281 relative_to_current = abs_path.parent / include_path
282 if relative_to_current.exists():
283 include_path = relative_to_current
284 else:
285 # Try relative to project root (parent of current file's parent)
286 # This handles cases like "snippets/datasets.yaml" from "snippets/features.yaml"
287 project_root = abs_path.parent.parent
288 relative_to_root = project_root / include_path
289 if relative_to_root.exists():
290 include_path = relative_to_root
291 else:
292 # Neither worked, raise error with both attempted paths
293 raise FileNotFoundError( # noqa: TRY003
294 f"Included file not found: {include_path_str}\n"
295 f" Tried: {relative_to_current}\n"
296 f" Tried: {relative_to_root}"
297 )
299 if not include_path.exists():
300 raise FileNotFoundError(f"Included file not found: {include_path}") # noqa: TRY003
302 # Recursively load the included file
303 included_data = self._load_with_includes(include_path, visited.copy())
305 # Merge included data into result (deep merge)
306 result = self._deep_merge(result, included_data)
308 # Merge current file's data on top (current file takes precedence)
309 return self._deep_merge(result, current_data)
311 def _deep_merge(
312 self, base: dict[str, Any], override: dict[str, Any]
313 ) -> dict[str, Any]:
314 """
315 Deep merge two dictionaries, with override taking precedence.
317 Parameters
318 ----------
319 base : Dict[str, Any]
320 Base dictionary
321 override : Dict[str, Any]
322 Override dictionary (takes precedence)
324 Returns
325 -------
326 Dict[str, Any]
327 Merged dictionary
328 """
329 result = base.copy()
331 for key, value in override.items():
332 if (
333 key in result
334 and isinstance(result[key], dict)
335 and isinstance(value, dict)
336 ):
337 # Recursively merge nested dictionaries
338 result[key] = self._deep_merge(result[key], value)
339 else:
340 # Override takes precedence
341 result[key] = value
343 return result
345 @classmethod
346 def from_file(
347 cls, yaml_path: str | Path, *, strict: bool = True, **variables: Any
348 ) -> YamlConfig:
349 return cls(yaml_path, strict=strict, **variables)
351 def to_dict(self) -> YamlDict:
352 return deepcopy(self._data)
354 def get_data(self) -> YamlDict:
355 """Get the raw configuration data dictionary.
357 Returns
358 -------
359 YamlDict
360 The raw configuration data dictionary
361 """
362 return self._data
364 def reload(self, *, strict: bool | None = None, **variables: Any) -> None:
365 """Reload the YAML file, optionally overriding variables and strictness.
367 Useful if the file or variables have changed.
368 """
370 new_strict = self.strict if strict is None else strict
371 updated_vars = {**self.variables, **variables}
372 refreshed = type(self).from_file(self.path, strict=new_strict, **updated_vars)
373 # Replace internal state
374 self._data = refreshed.get_data()
375 self.variables = refreshed.variables
376 self.strict = refreshed.strict
378 def get_datasets_list(self) -> list[str]:
379 """Return list of dataset names defined in the configuration.
381 Returns
382 -------
383 List[str]
384 List of dataset names, or empty list if no datasets section exists
385 """
386 datasets = self._data.get("datasets")
387 if not isinstance(datasets, dict):
388 return []
389 return list(datasets.keys())
391 def get_dataset_config(self, name: str) -> YamlDict:
392 """Return configuration dictionary for a specific dataset.
394 Parameters
395 ----------
396 name : str
397 Dataset name
399 Returns
400 -------
401 YamlDict
402 Dataset configuration as a plain dictionary
404 Raises
405 ------
406 KeyError
407 If dataset with given name does not exist
408 """
409 datasets = self._data.get("datasets")
410 if not isinstance(datasets, dict):
411 raise KeyError("No datasets section found in configuration") # noqa: TRY003
413 if name not in datasets:
414 raise KeyError(f"Dataset '{name}' not found in configuration") # noqa: TRY003
416 dataset_spec = datasets[name]
417 if isinstance(dataset_spec, dict):
418 return deepcopy(dataset_spec)
420 raise TypeError(f"Dataset '{name}' configuration is not a valid mapping") # noqa: TRY003
422 # Internal utilities
425__all__ = ["YamlConfig"]