From 6940e56e6c3c70f05336119f174689896a721650 Mon Sep 17 00:00:00 2001 From: Nathan Stender Date: Tue, 2 Jun 2026 11:29:08 -0400 Subject: [PATCH 1/2] fix: Upgrade jsonschema from <4.18.0 to >=4.18.0 The performance regression that caused the original pin (filed as referencing#178) has been resolved in newer versions. Migrate from the deprecated jsonschema.RefResolver API to the referencing.Registry API. Bundled sub-schemas with extended enums are correctly prioritized over standalone versions by comparing content size. Co-Authored-By: Claude Opus 4.6 --- pyproject.toml | 8 +--- src/allotropy/allotrope/schemas.py | 65 +++++++++++++++++++++--------- 2 files changed, 49 insertions(+), 24 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c6f7b0ccd..740b0c9f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,18 +50,14 @@ dependencies = [ "babel >= 2.9.0", "chardet >= 5.2.0", "defusedxml >= 0.7.1", - # NOTE: jsonschema 4.18.0 introduces a serious performance regression, seemingly due to use of new - # referencing library. - # Filed issue: https://github.com/python-jsonschema/referencing/issues/178 - # TODO(nstender): investigate removing all refs from schema before use to work around perf issues - # while unblocking package upgrades. - "jsonschema >= 4.3.3, < 4.18.0", + "jsonschema >= 4.18.0", "numpy >= 1.25.0", "olefile >= 0.47", "openpyxl >= 3.1.0", "pandas >= 2.2.0, < 3.0.0", "python-calamine >= 0.2.3", "rainbow-api == 1.0.10", + "referencing >= 0.31.0", "rfc3339-validator >= 0.1.4", "types-defusedxml >= 0.7.0.20240218", "xlrd >= 2.0.0", diff --git a/src/allotropy/allotrope/schemas.py b/src/allotropy/allotrope/schemas.py index fe7d7095f..58e3f7fe4 100644 --- a/src/allotropy/allotrope/schemas.py +++ b/src/allotropy/allotrope/schemas.py @@ -6,6 +6,8 @@ import jsonschema from jsonschema import ValidationError +from referencing import Registry, Resource +from referencing.jsonschema import DRAFT202012 from allotropy.allotrope.path_util import ( get_full_schema_path, @@ -225,24 +227,56 @@ def get_schema_from_model(model: Any) -> dict[str, Any]: return get_schema_from_manifest(manifest) -_schema_store: dict[str, dict[str, Any]] | None = None +_registry: Registry[dict[str, Any]] | None = None -def _get_schema_store() -> dict[str, dict[str, Any]]: - """Load all schemas into a dict keyed by their $id URI.""" - global _schema_store # noqa: PLW0603 - if _schema_store is not None: - return _schema_store +def _schema_content_size(schema: dict[str, Any]) -> int: + """Approximate content size of a schema's $defs for comparison.""" + defs = schema.get("$defs", {}) + return sum(len(json.dumps(v)) for v in defs.values()) if defs else 0 + + +def _get_registry() -> Registry[dict[str, Any]]: + """Build a referencing.Registry from all local schemas. + + Some technique schemas bundle extended copies of shared schemas (e.g. core.schema) + under $defs with their own $id. When a bundled copy is LARGER than the standalone + (more enum values, more definitions), we use the bundled version so that extended + schemas take precedence. When a bundled copy is a subset, we keep the standalone. + """ + global _registry # noqa: PLW0603 + if _registry is not None: + return _registry + + schemas_by_id: dict[str, dict[str, Any]] = {} + sizes_by_id: dict[str, int] = {} - store: dict[str, dict[str, Any]] = {} for schema_file in SCHEMA_DIR_PATH.rglob("*.schema.json"): with open(schema_file, encoding=DEFAULT_ENCODING) as f: schema = json.load(f) schema_id = schema.get("$id") - if schema_id: - store[schema_id] = schema - _schema_store = store - return store + if not schema_id: + continue + if schema_id not in schemas_by_id: + schemas_by_id[schema_id] = schema + sizes_by_id[schema_id] = _schema_content_size(schema) + + for val in schema.get("$defs", {}).values(): + if not isinstance(val, dict) or "$id" not in val: + continue + bundled_id = val["$id"] + bundled_size = _schema_content_size(val) + if bundled_size > sizes_by_id.get(bundled_id, -1): + schemas_by_id[bundled_id] = val + sizes_by_id[bundled_id] = bundled_size + + _registry = Registry().with_resources( + ( + (sid, Resource.from_contents(s, default_specification=DRAFT202012)) + for sid, s in schemas_by_id.items() + ) + ) + return _registry _schema_cache: dict[str, dict[str, Any]] = {} @@ -271,14 +305,9 @@ def _get_validator(schema_path: Path) -> Any: if cached is not None: return cached schema = _get_schema_by_path(schema_path) - store = _get_schema_store() - resolver = jsonschema.RefResolver( - base_uri=schema.get("$id", ""), - referrer=schema, - store=store, - ) + registry = _get_registry() validator = FastDraft202012Validator( - schema, resolver=resolver, format_checker=FORMAT_CHECKER + schema, registry=registry, format_checker=FORMAT_CHECKER ) _validator_cache[key] = validator return validator From 3353abbb64e630ee47113118df0964911c7d9ff6 Mon Sep 17 00:00:00 2001 From: Nathan Stender Date: Tue, 2 Jun 2026 12:00:30 -0400 Subject: [PATCH 2/2] perf: Add $ref lookup cache to eliminate redundant URL parsing A single validation triggers 40,000+ $ref resolutions with only ~17 unique URIs. Cache resolved schemas by (base_uri, ref) to avoid repeated URL parsing and pointer navigation. Brings performance to parity with jsonschema 4.17.3. Co-Authored-By: Claude Opus 4.6 --- src/allotropy/allotrope/schemas.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/src/allotropy/allotrope/schemas.py b/src/allotropy/allotrope/schemas.py index 58e3f7fe4..96853c3fb 100644 --- a/src/allotropy/allotrope/schemas.py +++ b/src/allotropy/allotrope/schemas.py @@ -197,9 +197,39 @@ def _fast_one_of_validator( yield from _original_one_of(validator, one_of, instance, schema) +# --------------------------------------------------------------------------- +# Cached $ref validator +# --------------------------------------------------------------------------- +# A single validation can trigger 40,000+ $ref lookups with only ~17 unique +# URIs. The referencing library's lookup() does URL parsing and pointer +# navigation on every call. We cache by (base_uri, ref) since the registry +# is a singleton and these two values fully determine the resolution result. + +_ref_lookup_cache: dict[tuple[str, str], tuple[Any, Any]] = {} + + +def _cached_ref_validator( + validator: Any, ref: str, instance: Any, schema: Any # noqa: ARG001 +) -> Any: + """$ref validator that caches resolver.lookup() results.""" + resolver = validator._resolver + cache_key = (resolver._base_uri, ref) + cached = _ref_lookup_cache.get(cache_key) + if cached is not None: + contents, resolved_resolver = cached + else: + resolved = resolver.lookup(ref) + contents = resolved.contents + resolved_resolver = resolved.resolver + _ref_lookup_cache[cache_key] = (contents, resolved_resolver) + + yield from validator.descend(instance, contents, resolver=resolved_resolver) + + FastDraft202012Validator = jsonschema.validators.extend( # type: ignore[no-untyped-call] jsonschema.validators.Draft202012Validator, validators={ + "$ref": _cached_ref_validator, "items": _fast_items_validator, "oneOf": _fast_one_of_validator, },