Coverage for databricks_ml_workbench/config.py: 71%

1from __future__ import annotations 

2 

3from copy import deepcopy 

4from pathlib import Path 

5from string import Formatter 

6from typing import Any, Dict, Iterable, List, Mapping, MutableMapping, Tuple, Union 

7 

8import yaml 

9 

10 

11YamlPrimitive = Union[str, int, float, bool, None] 

12YamlValue = Union[YamlPrimitive, "YamlDict", "YamlList"] 

13YamlDict = Dict[str, YamlValue] 

14YamlList = List[YamlValue] 

15 

16 

17def _extract_field_names(template: str) -> List[str]: 

18 """Return field names used in a ``str.format``-style template string. 

19 

20 Example: "{catalog}.{schema}.table" -> ["catalog", "schema"] 

21 """ 

22 

23 field_names: List[str] = [] 

24 for literal_text, field_name, format_spec, conversion in Formatter().parse(template): 

25 if field_name: # skip None and empty 

26 field_names.append(field_name) 

27 return field_names 

28 

29 

30def _interpolate_string(value: str, variables: Mapping[str, Any], *, strict: bool) -> str: 

31 """Interpolate placeholders in a string using ``variables``. 

32 

33 - If ``strict`` is True (default), raise a KeyError if any placeholder is missing. 

34 - If ``strict`` is False, leave unknown placeholders as-is. 

35 """ 

36 

37 if "{" not in value or "}" not in value: 

38 return value 

39 

40 placeholders = _extract_field_names(value) 

41 if not placeholders: 

42 return value 

43 

44 missing = [name for name in placeholders if name not in variables] 

45 if missing and strict: 

46 raise KeyError( 

47 f"Missing values for placeholders {missing} while formatting: {value!r}" 

48 ) 

49 

50 class _DefaultDict(dict): 

51 def __missing__(self, key: str) -> str: # type: ignore[override] 

52 return "{" + key + "}" 

53 

54 mapping: Mapping[str, Any] 

55 mapping = variables if strict else _DefaultDict(variables) # type: ignore[arg-type] 

56 return value.format_map(mapping) # type: ignore[arg-type] 

57 

58 

59def _interpolate_value(value: YamlValue, variables: Mapping[str, Any], *, strict: bool) -> YamlValue: 

60 if isinstance(value, str): 

61 return _interpolate_string(value, variables, strict=strict) 

62 if isinstance(value, dict): 

63 return {k: _interpolate_value(v, variables, strict=strict) for k, v in value.items()} 

64 if isinstance(value, list): 

65 return [_interpolate_value(v, variables, strict=strict) for v in value] 

66 return value 

67 

68 

69def _wrap(value: YamlValue) -> Any: 

70 if isinstance(value, dict): 

71 return _ConfigNode(value) 

72 if isinstance(value, list): 

73 return [_wrap(v) for v in value] 

74 return value 

75 

76 

77def _unwrap(value: Any) -> YamlValue: 

78 if isinstance(value, _ConfigNode): 

79 return {k: _unwrap(v) for k, v in value.items()} 

80 if isinstance(value, list): 

81 return [_unwrap(v) for v in value] 

82 return value 

83 

84 

85class _ConfigNode(MutableMapping[str, Any]): 

86 """Lightweight mapping that supports both dict- and attribute-style access. 

87 

88 This is used for nested configuration sections. 

89 """ 

90 

91 def __init__(self, data: YamlDict) -> None: 

92 self._data: YamlDict = data 

93 

94 # Mapping interface 

95 def __getitem__(self, key: str) -> Any: 

96 return _wrap(self._data[key]) 

97 

98 def __setitem__(self, key: str, value: Any) -> None: 

99 self._data[key] = _unwrap(value) 

100 

101 def __delitem__(self, key: str) -> None: 

102 del self._data[key] 

103 

104 def __iter__(self) -> Iterable[str]: 

105 return iter(self._data) 

106 

107 def __len__(self) -> int: # noqa: D401 - trivial 

108 return len(self._data) 

109 

110 # Attribute-style access 

111 def __getattr__(self, name: str) -> Any: 

112 try: 

113 return _wrap(self._data[name]) 

114 except KeyError as exc: 

115 raise AttributeError(name) from exc 

116 

117 def get(self, key: str, default: Any | None = None) -> Any: 

118 if key in self._data: 

119 return _wrap(self._data[key]) 

120 return default 

121 

122 def to_dict(self) -> YamlDict: 

123 return deepcopy(self._data) 

124 

125 def __repr__(self) -> str: # pragma: no cover - representation only 

126 return f"_ConfigNode({self._data!r})" 

127 

128 

129class YamlConfig(_ConfigNode): 

130 """Configuration loaded from a YAML file with placeholder interpolation. 

131 

132 Placeholders use Python's ``str.format`` style, e.g. "{catalog}". 

133 Values for placeholders are assembled as: 

134 1) ``defaults`` key from the YAML file (if present), then 

135 2) Keyword arguments provided to the constructor, which override defaults. 

136 

137 Parameters 

138 ---------- 

139 yaml_path: 

140 Path to the YAML file. 

141 strict: 

142 If True (default), raise when placeholders are missing. If False, leave 

143 unknown placeholders unchanged. 

144 **variables: 

145 Additional key/value pairs to use for placeholder interpolation. 

146 """ 

147 

148 def __init__(self, yaml_path: Union[str, Path], *, strict: bool = True, **variables: Any) -> None: 

149 path = Path(yaml_path) 

150 if not path.exists(): 

151 raise FileNotFoundError(f"YAML file not found: {path}") 

152 

153 with path.open("r", encoding="utf-8") as f: 

154 loaded: Any = yaml.safe_load(f) or {} 

155 

156 if not isinstance(loaded, dict): 

157 raise ValueError("Top-level YAML structure must be a mapping (dict)") 

158 

159 defaults_section: Mapping[str, Any] = loaded.get("defaults", {}) if isinstance(loaded.get("defaults", {}), dict) else {} 

160 

161 # Merge defaults with user-supplied variables (kwargs override defaults) 

162 variables_merged: Dict[str, Any] = {**defaults_section, **variables} 

163 

164 # Interpolate placeholders across the entire loaded YAML structure 

165 interpolated: YamlDict = _interpolate_value(loaded, variables_merged, strict=strict) # type: ignore[assignment] 

166 

167 super().__init__(interpolated) 

168 self.path: Path = path 

169 self.variables: Dict[str, Any] = variables_merged 

170 self.strict: bool = strict 

171 

172 @classmethod 

173 def from_file(cls, yaml_path: Union[str, Path], *, strict: bool = True, **variables: Any) -> "YamlConfig": 

174 return cls(yaml_path, strict=strict, **variables) 

175 

176 def to_dict(self) -> YamlDict: 

177 return deepcopy(self._data) 

178 

179 def reload(self, *, strict: bool | None = None, **variables: Any) -> None: 

180 """Reload the YAML file, optionally overriding variables and strictness. 

181 

182 Useful if the file or variables have changed. 

183 """ 

184 

185 new_strict = self.strict if strict is None else strict 

186 updated_vars = {**self.variables, **variables} 

187 refreshed = type(self).from_file(self.path, strict=new_strict, **updated_vars) 

188 # Replace internal state 

189 self._data = refreshed._data 

190 self.variables = refreshed.variables 

191 self.strict = refreshed.strict 

192 

193 

194__all__ = ["YamlConfig"] 

195 

196