Coverage for core / src / sensorkit / config / schema.py: 99%
92 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2from __future__ import annotations
4import importlib
5from collections.abc import Iterator
6from typing import Any
8from pydantic.json_schema import GenerateJsonSchema
10from sensorkit.config.parser import PARSER_VERSION, SensorKitBaseConfig
11from sensorkit.config.section import ConfigSection, config_sections
13SCHEMA_DIALECT = "https://json-schema.org/draft/2020-12/schema"
14SCHEMA_TITLE = "SensorKit configuration"
15ID_KEY_DESCRIPTION = "Name of the entity this entry defines."
17# Keywords that leave an object accepting properties beyond the ones it lists, either by
18# saying so outright or by composing with a schema that lists more.
19OPEN_KEYWORDS = frozenset(
20 {"additionalProperties", "allOf", "anyOf", "oneOf", "patternProperties", "$ref"}
21)
23type JsonSchema = dict[str, Any]
25# Key standing in for the base config in a definitions pass, where every other key is a
26# section name.
27BASE = None
30def _stitch_id_key(
31 section: ConfigSection,
32 schema: JsonSchema,
33 defs: dict[str, JsonSchema],
34) -> JsonSchema:
35 """Extend a section's schema with the entity ID key it accepts.
37 The key is merged into a copy of the model rather than composed onto a reference to it.
38 Composition would put the two property lists in separate subschemas, where neither one
39 can be closed against undeclared keys, and would leave the pair beyond the reach of the
40 JSON Schema support in common YAML editors.
42 Args:
43 section: The registered section the schema was generated for.
44 schema: The section's generated schema, either a reference to a model or a container
45 of them.
46 defs: The definitions the schema references.
48 Returns:
49 The section schema, extended where an ID key applies.
50 """
51 match section.id_source:
52 case "by_key":
53 slot = None
54 case "by_subkey":
55 slot = "items" if schema.get("type") == "array" else "additionalProperties"
56 case _:
57 # A section taking its IDs from mapping keys, or from an ID fixed at declaration,
58 # accepts no key of its own.
59 return schema
61 item = schema if slot is None else schema.get(slot)
62 ref = item.get("$ref") if isinstance(item, dict) else None
64 if ref is None:
65 return schema
67 model = defs[ref.rpartition("/")[2]]
68 declared = model.get("properties", {})
70 # Sections taking their IDs from a field of their own model need no stitching.
71 if section.id_key in declared:
72 return schema
74 extended: JsonSchema = {
75 **model,
76 "properties": {
77 section.id_key: {"type": "string", "description": ID_KEY_DESCRIPTION},
78 **declared,
79 },
80 }
82 if section.id_key_required:
83 extended["required"] = [section.id_key, *model.get("required", ())]
85 return extended if slot is None else {**schema, slot: extended}
88def _conditional_branches(prop: str, mapping: dict[str, str]) -> JsonSchema:
89 """Build a union that selects its branch by a discriminating property.
91 Args:
92 prop: Name of the property naming the branch to apply.
93 mapping: Reference to the branch each value of the property selects.
95 Returns:
96 An object schema applying exactly the branch the property names.
97 """
98 return {
99 "type": "object",
100 "properties": {prop: {"enum": list(mapping)}},
101 "required": [prop],
102 "allOf": [
103 {
104 "if": {"properties": {prop: {"const": value}}, "required": [prop]},
105 "then": {"$ref": ref},
106 }
107 for value, ref in mapping.items()
108 ],
109 }
112def _expand_discriminators(node: Any):
113 """Rewrite discriminated unions to select their branch conditionally, in place.
115 Pydantic marks a discriminated union with a `discriminator` annotation, an OpenAPI
116 extension that JSON Schema validators do not read. Left with the bare list of branches, a
117 validator tries all of them, and an entry matching none reports a single failure against
118 the whole object, which editors show against every key in it. Choosing the branch by its
119 discriminating property narrows the report to the keys responsible for it.
121 Args:
122 node: The schema document, or any node within it.
123 """
124 match node:
125 # The branches stay reachable as references, and the definitions they name are walked
126 # in their own right, so the rewritten node needs no further visiting.
127 case {"oneOf": _, "discriminator": {"propertyName": str(prop), "mapping": dict(mapping)}}:
128 node.clear()
129 node.update(_conditional_branches(prop, mapping))
130 case dict():
131 for value in node.values():
132 _expand_discriminators(value)
133 case list():
134 for item in node:
135 _expand_discriminators(item)
138def _closable(node: JsonSchema) -> bool:
139 """Report whether a node names every property it accepts.
141 Args:
142 node: The schema node to judge.
144 Returns:
145 True if the node's property list is exhaustive.
146 """
147 return node.get("type") == "object" and "properties" in node and not OPEN_KEYWORDS & node.keys()
150def _close_objects(node: Any):
151 """Forbid undeclared properties throughout a schema document, in place.
153 Pydantic describes a model as accepting more than its own fields only where the model
154 says so, so a model taking the default policy generates a property list an editor reads
155 as open. Closing the lists that are exhaustive turns a misspelled key into a reported
156 error rather than one that validation quietly drops.
158 Args:
159 node: The schema document, or any node within it.
160 """
161 match node:
162 case dict():
163 if _closable(node):
164 node["additionalProperties"] = False
166 for value in node.values():
167 _close_objects(value)
168 case list():
169 for item in node:
170 _close_objects(item)
173def _referenced_defs(node: Any) -> Iterator[str]:
174 """Yield the name of every definition referenced from a schema node.
176 Args:
177 node: The schema node to walk.
179 Yields:
180 Definition names, with repeats.
181 """
182 match node:
183 case dict():
184 if isinstance(ref := node.get("$ref"), str):
185 yield ref.rpartition("/")[2]
187 for value in node.values():
188 yield from _referenced_defs(value)
189 case list():
190 for item in node:
191 yield from _referenced_defs(item)
194def _reachable_defs(root: Any, defs: dict[str, JsonSchema]) -> dict[str, JsonSchema]:
195 """Select the definitions a document actually reaches.
197 Args:
198 root: The document to walk, excluding the definitions themselves.
199 defs: The generated definitions to select from.
201 Returns:
202 The definitions reachable from the document, in their generated order.
203 """
204 reachable: set[str] = set()
205 pending = list(_referenced_defs(root))
207 while pending:
208 name = pending.pop()
210 if name in reachable or name not in defs:
211 continue
213 reachable.add(name)
214 pending.extend(_referenced_defs(defs[name]))
216 return {name: schema for name, schema in defs.items() if name in reachable}
219def config_json_schema(*, title: str = SCHEMA_TITLE) -> JsonSchema:
220 """Generate a JSON Schema for a complete unified configuration file.
222 The schema covers the fixed top-level keys along with every config section registered so
223 far. Sections come from modules, so import the modules a site uses before calling this
224 (`import_modules` does so from the environment and the config file itself), or the
225 result describes only the core sections.
227 Args:
228 title: Title recorded in the generated schema.
230 Returns:
231 A JSON Schema document, ready to serialize.
232 """
233 # Make sure core sections are imported.
234 importlib.import_module("sensorkit.config.core")
236 sections = sorted(config_sections().items())
237 generator = GenerateJsonSchema()
239 # Generating everything in one pass keeps the definitions unified and lets Pydantic
240 # disambiguate models that share a name across modules.
241 refs, defs = generator.generate_definitions(
242 [
243 (BASE, "validation", SensorKitBaseConfig.__pydantic_core_schema__),
244 *((key, "validation", section.adapter.core_schema) for key, section in sections),
245 ]
246 )
248 # The base config becomes the root of the document, so its definition is consumed rather
249 # than referenced.
250 base = defs.pop(refs[(BASE, "validation")]["$ref"].rpartition("/")[2])
251 properties = dict(base.get("properties", {}))
253 # The parser accepts only its own version, which the field default does not convey.
254 properties["version"] = {**properties["version"], "const": PARSER_VERSION}
255 properties["version"].pop("default", None)
257 properties.update(
258 {
259 key: _stitch_id_key(section, refs[(key, "validation")], defs)
260 for key, section in sections
261 }
262 )
264 document = {
265 "$schema": SCHEMA_DIALECT,
266 "title": title,
267 "type": "object",
268 "properties": properties,
269 "required": ["version"],
270 # An unregistered top-level key is a config error, not free-form data.
271 "additionalProperties": False,
272 # Stitching an ID key consumes the section's own definition.
273 "$defs": _reachable_defs(properties, defs),
274 }
276 _expand_discriminators(document)
277 _close_objects(document)
279 return document