Coverage for core / src / sensorkit / config / parser.py: 80%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import collections 

5import importlib 

6import itertools 

7from collections.abc import Iterable, Mapping 

8from dataclasses import dataclass 

9from typing import Any 

10 

11from loguru import logger 

12from pydantic import BaseModel, Field, model_validator 

13 

14from sensorkit.common.keyword import KeywordDict 

15from sensorkit.config.section import ConfigSection, get_config_section 

16 

17PARSER_VERSION = 1 

18 

19 

20class GeneralConfig(BaseModel, extra="forbid"): 

21 backend: str | None = None 

22 modules: list[str] = Field(default_factory=list) 

23 imports: list[str] = Field(default_factory=list) 

24 

25 

26class ServiceConfig(BaseModel, extra="forbid"): 

27 id: str 

28 python_path: str | None = None 

29 python_module: str | None = None 

30 python_file: str | None = None 

31 

32 @model_validator(mode="after") 

33 def _validator(self): 

34 if bool(self.python_module) == bool(self.python_path): 

35 raise ValueError("Exactly one of 'python_module' or 'python_path' must be set") 

36 

37 if self.python_file and self.python_path: 

38 raise ValueError("Cannot set both 'python_file' and 'python_path'") 

39 

40 return self 

41 

42 

43class SensorKitBaseConfig(BaseModel, extra="allow"): 

44 version: int = 0 

45 sensorkit: GeneralConfig = Field(default_factory=GeneralConfig) 

46 globals: KeywordDict = Field(default_factory=KeywordDict) 

47 services: list[ServiceConfig] = Field(default_factory=list) 

48 

49 @model_validator(mode="after") 

50 def _validate(self): 

51 if self.version != PARSER_VERSION: 

52 raise ConfigVersionUnsupported(f"Unsupported config version: {self.version}") 

53 

54 return self 

55 

56 def resolve_dynamic_sections(self): 

57 # Make sure core sections are imported. 

58 importlib.import_module("sensorkit.config.core") 

59 

60 ekv, services = _parse_extra_sections(self.model_extra) 

61 return SensorKitConfig( 

62 base=self, 

63 services=self.services + services, 

64 entity_kv=ekv, 

65 ) 

66 

67 def configured_imports(self) -> Iterable[str]: 

68 # Modules must come first, in case any imported user code uses them. 

69 return itertools.chain( 

70 (f"sensorkit.{module}" for module in self.sensorkit.modules), 

71 self.sensorkit.imports, 

72 ) 

73 

74 

75@dataclass 

76class SensorKitConfig: 

77 base: SensorKitBaseConfig 

78 services: list[ServiceConfig] 

79 entity_kv: dict[str, list[BaseModel]] 

80 

81 @property 

82 def backend(self): 

83 return self.base.sensorkit.backend 

84 

85 @property 

86 def global_kv(self): 

87 return self.base.globals 

88 

89 

90def _validate_mapper_return_value[T](value: Any, expected_type: type[T]) -> tuple[T] | None: 

91 if isinstance(value, expected_type): 

92 return (value,) 

93 elif isinstance(value, Iterable): 

94 value = tuple(value) 

95 

96 if not isinstance(value, tuple) or not all(isinstance(elem, expected_type) for elem in value): 

97 logger.debug(f"Invalid mapper return value for type {expected_type.__name__}: {value}") 

98 return None 

99 

100 return value 

101 

102 

103def _parse_section(section: ConfigSection, value: Any) -> Iterable[tuple[str, BaseModel]]: 

104 try: 

105 ids = _validate_mapper_return_value(section.id_mapper(value), str) 

106 except Exception as e: 

107 raise ConfigError(f"Error in config section {section.key!r} (id mapper error)") from e 

108 

109 if ids is None: 

110 raise ConfigError(f"Error in config section {section.key!r} (bad id mapper?)") 

111 

112 instance = section.adapter.validate_python(value) 

113 

114 try: 

115 models = _validate_mapper_return_value( 

116 section.model_mapper(instance) if section.model_mapper is not None else instance, 

117 BaseModel, 

118 ) 

119 except Exception as e: 

120 raise ConfigError(f"Error in config section {section.key!r} (model mapper error)") from e 

121 

122 if models is None: 

123 raise ConfigError(f"Error in config section {section.key!r} (bad model mapper?)") 

124 

125 return zip(ids, models, strict=True) 

126 

127 

128def _parse_extra_sections(config: Mapping[str, Any]): 

129 ekv: dict[str, list[BaseModel]] = collections.defaultdict(list) 

130 services: list[ServiceConfig] = [] 

131 

132 for key, value in config.items(): 

133 section = get_config_section(key) 

134 

135 if not section: 

136 raise ConfigSectionUnknown(f"Unknown config section {key!r}") 

137 

138 for entity_id, model in _parse_section(section, value): 

139 ekv[entity_id].append(model) 

140 

141 if section.service_path: 

142 services.append(ServiceConfig(id=entity_id, python_module=section.service_path)) 

143 

144 return ekv, services 

145 

146 

147def parse_config(config: Mapping[str, Any]) -> SensorKitBaseConfig: 

148 """Parse the given unified configuration dict. 

149 

150 Args: 

151 config: A mapping containing the unified configuration data. 

152 

153 Returns: 

154 A SensorKitBaseConfig instance. 

155 

156 Raises: 

157 ConfigVersionUnsupported: The config version is mismatched 

158 ConfigSectionUnknown: A config section cannot be found in the registry 

159 ConfigError: An internal error with config processing 

160 ValidationError: A config section fails Pydantic validation during parsing 

161 """ 

162 return SensorKitBaseConfig.model_validate(config) 

163 

164 

165class ConfigError(Exception): 

166 """Raised when an error occurs during config loading.""" 

167 

168 

169class ConfigVersionUnsupported(ConfigError): 

170 """Raised when the config version is unsupported.""" 

171 

172 

173class ConfigSectionUnknown(ConfigError): 

174 """Raised when the config section is unknown."""