Coverage for core / src / sensorkit / config / section.py: 100%

1# SPDX-License-Identifier: Apache-2.0 

2import operator 

3from collections.abc import Iterable, Mapping 

4from typing import Any, Callable, Literal, NamedTuple 

5 

6from pydantic import BaseModel, TypeAdapter 

7 

8type UnifiedConfigIdMapper = Callable[[Any], Iterable[str]] 

9type UnifiedConfigModelMapper[T] = Callable[[T], Iterable[BaseModel]] 

10type IdSource = Literal["by_key", "by_subkey", "mapping_key", "default"] 

11 

12DEFAULT_ID_KEY = "id" 

13KEYED_ID_SOURCES = ("by_key", "by_subkey") 

14DEFAULTABLE_ID_SOURCES = ("by_key", "default") 

15 

16 

17def _section_entries(value: Any) -> Iterable[Any]: 

18 return value.values() if isinstance(value, Mapping) else value 

19 

20 

21def _validate_id_naming(key: str, id_source: str, id_key: str | None, id_default: str | None): 

22 """Reject a declaration whose ID arguments contradict the source it names. 

23 

24 Args: 

25 key: Top-level YAML key the section claims. 

26 id_source: Where the section takes its entity IDs from. 

27 id_key: Key holding the entity ID, where one was given. 

28 id_default: Entity ID to fall back on, where one was given. 

29 

30 Raises: 

31 ValueError: The arguments cannot apply together. 

32 """ 

33 if id_key is not None and id_source not in KEYED_ID_SOURCES: 

34 raise ValueError(f"Config section {key!r} cannot take 'id_key' with {id_source!r}") 

35 

36 # One ID cannot stand in for entries that omit their key, which would leave several 

37 # entries sharing an entity. 

38 if id_default is not None and id_source not in DEFAULTABLE_ID_SOURCES: 

39 raise ValueError(f"Config section {key!r} cannot take 'id_default' with {id_source!r}") 

40 

41 if id_source == "default" and id_default is None: 

42 raise ValueError(f"Config section {key!r} must give an 'id_default' with {id_source!r}") 

43 

44 

45class ConfigSection(NamedTuple): 

46 """A config section handler. 

47 

48 Args: 

49 key: Top-level YAML key this handler claims. 

50 adapter: Validates the section's raw value. 

51 id_source: Where the section takes its entity IDs from. 

52 id_mapper: Takes the raw value and returns the entity ID to write each model under, 

53 one per model the model mapper yields. 

54 model_mapper: Takes the validated instance and yields the models to write to KV, one 

55 per ID. None writes the validated instance itself. 

56 service_path: Path to a service implementation, where the section launches one. 

57 id_key: Key holding the entity ID, for the keyed ID sources. 

58 id_key_required: Whether the file has to supply `id_key`. False where the section 

59 falls back to an ID of its own. 

60 """ 

61 

62 key: str 

63 adapter: TypeAdapter 

64 id_source: IdSource 

65 id_mapper: UnifiedConfigIdMapper 

66 model_mapper: UnifiedConfigModelMapper | None = None 

67 service_path: str | None = None 

68 id_key: str | None = None 

69 id_key_required: bool = False 

70 

71 

72_registry: dict[str, ConfigSection] = {} 

73 

74 

75def declare_config_section[T]( 

76 key: str, 

77 parse_type: Any, 

78 *, 

79 id_source: IdSource, 

80 id_key: str | None = None, 

81 id_default: str | None = None, 

82 id_mapper: UnifiedConfigIdMapper | None = None, 

83 model_mapper: UnifiedConfigModelMapper[T] | None = None, 

84 service_path: str | None = None, 

85): 

86 """Declare and register a config section handler. 

87 

88 The section's records are validated against `parse_type` and written to the KV namespace 

89 of the entity each one configures. Where those IDs come from is declared rather than 

90 computed, so the generated JSON Schema can describe the key carrying them alongside the 

91 model's own fields. 

92 

93 Args: 

94 key: Top-level YAML key this handler claims. 

95 parse_type: Pydantic-compatible type used to validate the raw YAML value. Accepts any 

96 form TypeAdapter understands: type[T], list[T], etc. 

97 id_source: Where the section takes the entity ID for each of its records. `by_key` 

98 reads it from a key of the value itself, for a section configuring one entity. 

99 `by_subkey` reads it from a key of each entry, for a section configuring one 

100 entity per entry. `mapping_key` takes the IDs from the value's own keys. 

101 `default` names the entity at declaration, and the file supplies no ID at all. 

102 id_key: Key holding the entity ID, for `by_key` and `by_subkey`. Defaults to `id`. 

103 id_default: Entity ID to use where the file supplies none. Required with `default`, 

104 which admits no other ID. Given with `by_key` it makes the key optional. 

105 id_mapper: Optional callable that takes the raw value and returns an entity ID per 

106 model, replacing the one `id_source` implies. Sections producing several models 

107 from a single entry need one, since IDs are matched against the models rather 

108 than against the entries. 

109 model_mapper: Optional callable that takes the validated instance and yields the 

110 models to write to KV, one per ID. Defaults to the section's entries. 

111 service_path: Optional path to a service implementation. 

112 

113 Raises: 

114 ValueError: The section is already registered, or its ID naming is inconsistent. 

115 """ 

116 _validate_id_naming(key, id_source, id_key, id_default) 

117 

118 resolved_key = id_key or DEFAULT_ID_KEY 

119 

120 match id_source: 

121 case "by_key": 

122 derived = ( 

123 operator.itemgetter(resolved_key) 

124 if id_default is None 

125 else lambda value: value.get(resolved_key, id_default) 

126 ) 

127 case "by_subkey": 

128 derived = lambda value: (entry[resolved_key] for entry in _section_entries(value)) 

129 model_mapper = model_mapper or _section_entries 

130 case "mapping_key": 

131 derived = iter 

132 model_mapper = model_mapper or _section_entries 

133 case "default": 

134 derived = lambda _: id_default 

135 case _: 

136 raise ValueError(f"Config section {key!r} has an unknown ID source {id_source!r}") 

137 

138 if key in _registry: 

139 raise ValueError(f"Config section {key!r} already registered") 

140 

141 keyed = id_source in KEYED_ID_SOURCES 

142 

143 _registry[key] = ConfigSection( 

144 key=key, 

145 adapter=TypeAdapter(parse_type), 

146 id_source=id_source, 

147 id_mapper=id_mapper or derived, 

148 model_mapper=model_mapper, 

149 service_path=service_path, 

150 id_key=resolved_key if keyed else None, 

151 id_key_required=keyed and id_default is None, 

152 ) 

153 

154 

155def get_config_section(key: str) -> ConfigSection | None: 

156 """Retrieve a registered config section handler. 

157 

158 Args: 

159 key: The top-level YAML key to look up. 

160 

161 Returns: 

162 The registered section, or None if the key is not found in the registry. 

163 """ 

164 return _registry.get(key) 

165 

166 

167def config_sections() -> dict[str, ConfigSection]: 

168 """Retrieve every registered config section handler. 

169 

170 Returns: 

171 A snapshot of the registry, keyed by top-level YAML key. 

172 """ 

173 return dict(_registry)