chore: Revert "fix: Upgrade jsonschema from <4.18.0 to >=4.18.0 (#1217)"

pyproject.toml

@@ -50,14 +50,18 @@ dependencies = [
   "babel >= 2.9.0",
   "chardet >= 5.2.0",
   "defusedxml >= 0.7.1",
-  "jsonschema >= 4.18.0",
+  # 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",
   "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",

src/allotropy/allotrope/schemas.py

@@ -6,8 +6,6 @@
 
 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,
@@ -197,39 +195,9 @@ 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,
     },
@@ -257,67 +225,24 @@ def get_schema_from_model(model: Any) -> dict[str, Any]:
     return get_schema_from_manifest(manifest)
 
 
-_registry: Registry[dict[str, Any]] | None = None
-
-
-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.
+_schema_store: dict[str, dict[str, Any]] | None = None
 
-    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.
 
-    Resolution is fully deterministic: files are sorted, the largest version of each $id
-    wins, and on size ties standalone schemas take precedence over bundled copies.
-    """
-    global _registry  # noqa: PLW0603
-    if _registry is not None:
-        return _registry
+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
 
-    schemas_by_id: dict[str, dict[str, Any]] = {}
-    sizes_by_id: dict[str, int] = {}
-    is_standalone: dict[str, bool] = {}
-
-    for schema_file in sorted(SCHEMA_DIR_PATH.rglob("*.schema.json")):
+    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 not schema_id:
-            continue
-        schema_size = _schema_content_size(schema)
-        existing_size = sizes_by_id.get(schema_id, -1)
-        if schema_size > existing_size or (
-            schema_size == existing_size and not is_standalone.get(schema_id, False)
-        ):
-            schemas_by_id[schema_id] = schema
-            sizes_by_id[schema_id] = schema_size
-            is_standalone[schema_id] = True
-
-        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)
-            existing_size = sizes_by_id.get(bundled_id, -1)
-            if bundled_size > existing_size:
-                schemas_by_id[bundled_id] = val
-                sizes_by_id[bundled_id] = bundled_size
-                is_standalone[bundled_id] = False
-
-    _registry = Registry().with_resources(
-        (
-            (sid, Resource.from_contents(s, default_specification=DRAFT202012))
-            for sid, s in schemas_by_id.items()
-        )
-    )
-    return _registry
+        if schema_id:
+            store[schema_id] = schema
+    _schema_store = store
+    return store
 
 
 _schema_cache: dict[str, dict[str, Any]] = {}
@@ -346,9 +271,14 @@ def _get_validator(schema_path: Path) -> Any:
     if cached is not None:
         return cached
     schema = _get_schema_by_path(schema_path)
-    registry = _get_registry()
+    store = _get_schema_store()
+    resolver = jsonschema.RefResolver(
+        base_uri=schema.get("$id", ""),
+        referrer=schema,
+        store=store,
+    )
     validator = FastDraft202012Validator(
-        schema, registry=registry, format_checker=FORMAT_CHECKER
+        schema, resolver=resolver, format_checker=FORMAT_CHECKER
     )
     _validator_cache[key] = validator
     return validator