|
6 | 6 |
|
7 | 7 | import jsonschema |
8 | 8 | from jsonschema import ValidationError |
| 9 | +from referencing import Registry, Resource |
| 10 | +from referencing.jsonschema import DRAFT202012 |
9 | 11 |
|
10 | 12 | from allotropy.allotrope.path_util import ( |
11 | 13 | get_full_schema_path, |
@@ -195,9 +197,39 @@ def _fast_one_of_validator( |
195 | 197 | yield from _original_one_of(validator, one_of, instance, schema) |
196 | 198 |
|
197 | 199 |
|
| 200 | +# --------------------------------------------------------------------------- |
| 201 | +# Cached $ref validator |
| 202 | +# --------------------------------------------------------------------------- |
| 203 | +# A single validation can trigger 40,000+ $ref lookups with only ~17 unique |
| 204 | +# URIs. The referencing library's lookup() does URL parsing and pointer |
| 205 | +# navigation on every call. We cache by (base_uri, ref) since the registry |
| 206 | +# is a singleton and these two values fully determine the resolution result. |
| 207 | + |
| 208 | +_ref_lookup_cache: dict[tuple[str, str], tuple[Any, Any]] = {} |
| 209 | + |
| 210 | + |
| 211 | +def _cached_ref_validator( |
| 212 | + validator: Any, ref: str, instance: Any, schema: Any # noqa: ARG001 |
| 213 | +) -> Any: |
| 214 | + """$ref validator that caches resolver.lookup() results.""" |
| 215 | + resolver = validator._resolver |
| 216 | + cache_key = (resolver._base_uri, ref) |
| 217 | + cached = _ref_lookup_cache.get(cache_key) |
| 218 | + if cached is not None: |
| 219 | + contents, resolved_resolver = cached |
| 220 | + else: |
| 221 | + resolved = resolver.lookup(ref) |
| 222 | + contents = resolved.contents |
| 223 | + resolved_resolver = resolved.resolver |
| 224 | + _ref_lookup_cache[cache_key] = (contents, resolved_resolver) |
| 225 | + |
| 226 | + yield from validator.descend(instance, contents, resolver=resolved_resolver) |
| 227 | + |
| 228 | + |
198 | 229 | FastDraft202012Validator = jsonschema.validators.extend( # type: ignore[no-untyped-call] |
199 | 230 | jsonschema.validators.Draft202012Validator, |
200 | 231 | validators={ |
| 232 | + "$ref": _cached_ref_validator, |
201 | 233 | "items": _fast_items_validator, |
202 | 234 | "oneOf": _fast_one_of_validator, |
203 | 235 | }, |
@@ -225,24 +257,56 @@ def get_schema_from_model(model: Any) -> dict[str, Any]: |
225 | 257 | return get_schema_from_manifest(manifest) |
226 | 258 |
|
227 | 259 |
|
228 | | -_schema_store: dict[str, dict[str, Any]] | None = None |
| 260 | +_registry: Registry[dict[str, Any]] | None = None |
| 261 | + |
| 262 | + |
| 263 | +def _schema_content_size(schema: dict[str, Any]) -> int: |
| 264 | + """Approximate content size of a schema's $defs for comparison.""" |
| 265 | + defs = schema.get("$defs", {}) |
| 266 | + return sum(len(json.dumps(v)) for v in defs.values()) if defs else 0 |
229 | 267 |
|
230 | 268 |
|
231 | | -def _get_schema_store() -> dict[str, dict[str, Any]]: |
232 | | - """Load all schemas into a dict keyed by their $id URI.""" |
233 | | - global _schema_store # noqa: PLW0603 |
234 | | - if _schema_store is not None: |
235 | | - return _schema_store |
| 269 | +def _get_registry() -> Registry[dict[str, Any]]: |
| 270 | + """Build a referencing.Registry from all local schemas. |
| 271 | +
|
| 272 | + Some technique schemas bundle extended copies of shared schemas (e.g. core.schema) |
| 273 | + under $defs with their own $id. When a bundled copy is LARGER than the standalone |
| 274 | + (more enum values, more definitions), we use the bundled version so that extended |
| 275 | + schemas take precedence. When a bundled copy is a subset, we keep the standalone. |
| 276 | + """ |
| 277 | + global _registry # noqa: PLW0603 |
| 278 | + if _registry is not None: |
| 279 | + return _registry |
| 280 | + |
| 281 | + schemas_by_id: dict[str, dict[str, Any]] = {} |
| 282 | + sizes_by_id: dict[str, int] = {} |
236 | 283 |
|
237 | | - store: dict[str, dict[str, Any]] = {} |
238 | 284 | for schema_file in SCHEMA_DIR_PATH.rglob("*.schema.json"): |
239 | 285 | with open(schema_file, encoding=DEFAULT_ENCODING) as f: |
240 | 286 | schema = json.load(f) |
241 | 287 | schema_id = schema.get("$id") |
242 | | - if schema_id: |
243 | | - store[schema_id] = schema |
244 | | - _schema_store = store |
245 | | - return store |
| 288 | + if not schema_id: |
| 289 | + continue |
| 290 | + if schema_id not in schemas_by_id: |
| 291 | + schemas_by_id[schema_id] = schema |
| 292 | + sizes_by_id[schema_id] = _schema_content_size(schema) |
| 293 | + |
| 294 | + for val in schema.get("$defs", {}).values(): |
| 295 | + if not isinstance(val, dict) or "$id" not in val: |
| 296 | + continue |
| 297 | + bundled_id = val["$id"] |
| 298 | + bundled_size = _schema_content_size(val) |
| 299 | + if bundled_size > sizes_by_id.get(bundled_id, -1): |
| 300 | + schemas_by_id[bundled_id] = val |
| 301 | + sizes_by_id[bundled_id] = bundled_size |
| 302 | + |
| 303 | + _registry = Registry().with_resources( |
| 304 | + ( |
| 305 | + (sid, Resource.from_contents(s, default_specification=DRAFT202012)) |
| 306 | + for sid, s in schemas_by_id.items() |
| 307 | + ) |
| 308 | + ) |
| 309 | + return _registry |
246 | 310 |
|
247 | 311 |
|
248 | 312 | _schema_cache: dict[str, dict[str, Any]] = {} |
@@ -271,14 +335,9 @@ def _get_validator(schema_path: Path) -> Any: |
271 | 335 | if cached is not None: |
272 | 336 | return cached |
273 | 337 | schema = _get_schema_by_path(schema_path) |
274 | | - store = _get_schema_store() |
275 | | - resolver = jsonschema.RefResolver( |
276 | | - base_uri=schema.get("$id", ""), |
277 | | - referrer=schema, |
278 | | - store=store, |
279 | | - ) |
| 338 | + registry = _get_registry() |
280 | 339 | validator = FastDraft202012Validator( |
281 | | - schema, resolver=resolver, format_checker=FORMAT_CHECKER |
| 340 | + schema, registry=registry, format_checker=FORMAT_CHECKER |
282 | 341 | ) |
283 | 342 | _validator_cache[key] = validator |
284 | 343 | return validator |
|
0 commit comments