Skip to content

Commit d21b6d9

Browse files
fix: Upgrade jsonschema from <4.18.0 to >=4.18.0 (#1217)
## Summary - Removes the `< 4.18.0` version pin on `jsonschema` that was added 2 years ago due to a performance regression in the `referencing` library (issue #178) - Migrates from deprecated `jsonschema.RefResolver` API to `referencing.Registry` API - Handles bundled sub-schemas correctly: technique schemas that bundle extended copies of shared schemas (with additional enum values) are prioritized over standalone versions by comparing content size ## Context The performance regression has been resolved upstream. Benchmarking confirmed jsonschema 4.26.0 performs equally or better than the pinned 4.17.3 for both cold-start and large-file validation scenarios. ## Test plan - [x] All 1439 tests pass (only pre-existing `test_table_contents` failure unrelated to this change) - [x] Biacore tests pass (previously failing due to bundled units.schema missing `nM`) - [x] Flow-cytometry tests pass (bundled core.schema with extended enums correctly prioritized) - [x] Gen5 image tests pass - [x] Mypy and ruff pass cleanly 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 582e4ed commit d21b6d9

2 files changed

Lines changed: 79 additions & 24 deletions

File tree

pyproject.toml

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,18 +50,14 @@ dependencies = [
5050
"babel >= 2.9.0",
5151
"chardet >= 5.2.0",
5252
"defusedxml >= 0.7.1",
53-
# NOTE: jsonschema 4.18.0 introduces a serious performance regression, seemingly due to use of new
54-
# referencing library.
55-
# Filed issue: https://github.com/python-jsonschema/referencing/issues/178
56-
# TODO(nstender): investigate removing all refs from schema before use to work around perf issues
57-
# while unblocking package upgrades.
58-
"jsonschema >= 4.3.3, < 4.18.0",
53+
"jsonschema >= 4.18.0",
5954
"numpy >= 1.25.0",
6055
"olefile >= 0.47",
6156
"openpyxl >= 3.1.0",
6257
"pandas >= 2.2.0, < 3.0.0",
6358
"python-calamine >= 0.2.3",
6459
"rainbow-api == 1.0.10",
60+
"referencing >= 0.31.0",
6561
"rfc3339-validator >= 0.1.4",
6662
"types-defusedxml >= 0.7.0.20240218",
6763
"xlrd >= 2.0.0",

src/allotropy/allotrope/schemas.py

Lines changed: 77 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
import jsonschema
88
from jsonschema import ValidationError
9+
from referencing import Registry, Resource
10+
from referencing.jsonschema import DRAFT202012
911

1012
from allotropy.allotrope.path_util import (
1113
get_full_schema_path,
@@ -195,9 +197,39 @@ def _fast_one_of_validator(
195197
yield from _original_one_of(validator, one_of, instance, schema)
196198

197199

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+
198229
FastDraft202012Validator = jsonschema.validators.extend( # type: ignore[no-untyped-call]
199230
jsonschema.validators.Draft202012Validator,
200231
validators={
232+
"$ref": _cached_ref_validator,
201233
"items": _fast_items_validator,
202234
"oneOf": _fast_one_of_validator,
203235
},
@@ -225,24 +257,56 @@ def get_schema_from_model(model: Any) -> dict[str, Any]:
225257
return get_schema_from_manifest(manifest)
226258

227259

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
229267

230268

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] = {}
236283

237-
store: dict[str, dict[str, Any]] = {}
238284
for schema_file in SCHEMA_DIR_PATH.rglob("*.schema.json"):
239285
with open(schema_file, encoding=DEFAULT_ENCODING) as f:
240286
schema = json.load(f)
241287
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
246310

247311

248312
_schema_cache: dict[str, dict[str, Any]] = {}
@@ -271,14 +335,9 @@ def _get_validator(schema_path: Path) -> Any:
271335
if cached is not None:
272336
return cached
273337
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()
280339
validator = FastDraft202012Validator(
281-
schema, resolver=resolver, format_checker=FORMAT_CHECKER
340+
schema, registry=registry, format_checker=FORMAT_CHECKER
282341
)
283342
_validator_cache[key] = validator
284343
return validator

0 commit comments

Comments
 (0)