From 520e04a43957d8a522735a967b47020ed4884f07 Mon Sep 17 00:00:00 2001
From: "Ed (ODSC)" <ed.lloyd-davies@opendataservices.coop>
Date: Tue, 30 Apr 2024 15:33:14 +0100
Subject: [PATCH] sphinxcontrib/jsonschema.py: Add urn handling

https://github.com/openownership/data-standard/issues/546
---
 .github/workflows/test.yml  |  2 +-
 requirements_dev.txt        |  3 ++-
 setup.py                    |  3 +++
 sphinxcontrib/jsonschema.py | 45 +++++++++++++++++++++++++++++++++----
 4 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 636787a..0d435f6 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -8,7 +8,7 @@ jobs:
     strategy:
       matrix:
         python-version: [  '3.8', '3.9', '3.10', '3.11']
-        myst-parser-version: [ '<0.18.0', '>=0.18.0,<0.19', '>=0.19.0,<1.0', '>=1.0.0,<2', '>=2.0.0,<3']
+        myst-parser-version: [ '>=0.18.0,<0.19', '>=0.19.0,<1.0', '>=1.0.0,<2', '>=2.0.0,<3']
         jsonref-version: [">1"]
         include:
           # jsonref 1.0 has a backwards incompatible change - make sure we test just once with an older version of jsonref
diff --git a/requirements_dev.txt b/requirements_dev.txt
index 52980a2..22a3555 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -1,3 +1,4 @@
 -e .
 sphinx
-flake8
+lxml
+flake8<6
diff --git a/setup.py b/setup.py
index c9e79a2..84c0959 100644
--- a/setup.py
+++ b/setup.py
@@ -38,11 +38,14 @@
         'jsonref',
         'jsonpointer',
         'myst-parser',
+        'referencing',
+        'jscc',
     ],
     extras_require={
         'test': [
             'flake8<6',
             'lxml',
+            'defusedxml',
             'pytest',
         ],
     },
diff --git a/sphinxcontrib/jsonschema.py b/sphinxcontrib/jsonschema.py
index 1273e1b..a149186 100644
--- a/sphinxcontrib/jsonschema.py
+++ b/sphinxcontrib/jsonschema.py
@@ -13,6 +13,11 @@
 from docutils import nodes
 from docutils.parsers.rst import directives, Directive
 from pathlib import Path
+from urllib import parse as urlparse
+from referencing import Registry, Resource
+from referencing.jsonschema import DRAFT202012
+from jscc.schema import is_json_schema
+from jscc.testing.filesystem import walk_json_data
 
 import json
 from collections import OrderedDict
@@ -38,6 +43,26 @@ def custom_jsonref_jsonloader(uri, **kwargs):
     return {}
 
 
+def build_custom_schema_loader(schema_path):
+    """
+    Buildss a callable which handles a URN if provided (e.g. resolves part
+    before hash in 'urn:components#/$defs/UnspecifiedRecord' to components.json
+    in schema_path dircetory), else calls default JSON loader.
+    """
+    schemas = []
+    for _, _, _, data in walk_json_data(top=schema_path):
+        if is_json_schema(data):
+            schemas.append((data.get("$id"), Resource(contents=data, specification=DRAFT202012)))
+    registry = Registry().with_resources(schemas)
+    def custom_loader(uri, **kwargs):
+        scheme = urlparse.urlsplit(uri).scheme
+        if scheme == "urn":
+            return registry.contents(uri.split("#")[0])
+        else:
+            return jsonref.jsonloader(uri, **kwargs)
+    return custom_loader
+
+
 class JSONSchemaDirective(Directive):
     has_content = True
     required_arguments = 1
@@ -49,6 +74,7 @@ class JSONSchemaDirective(Directive):
         'addtargets': directives.flag,
         'externallinks': directives.unchanged,
         'allowexternalrefs': directives.flag,
+        'allowurnrefs': directives.flag,
     }
     # Add a rollup option here
 
@@ -87,11 +113,18 @@ def run(self):
                                        self.arguments[0])
                 env.note_dependency(relpath)
 
-                schema = JSONSchema.loadfromfile(abspath, allow_external_refs=('allowexternalrefs' in self.options))
+                schema = JSONSchema.loadfromfile(
+                    abspath,
+                    allow_external_refs=('allowexternalrefs' in self.options),
+                    loader=(build_custom_schema_loader(abspath.rsplit("/", 1)[0])
+                            if 'allowurnrefs' in self.options else None)
+                )
             else:
                 schema = JSONSchema.loadfromfile(
                     ''.join(self.content),
-                    allow_external_refs=('allowexternalrefs' in self.options)
+                    allow_external_refs=('allowexternalrefs' in self.options),
+                    loader=(build_custom_schema_loader(abspath.rsplit("/", 1)[0])
+                            if 'allowurnrefs' in self.options else None)
                 )
         except ValueError as exc:
             raise self.error('Failed to parse JSON Schema: %s' % exc)
@@ -242,10 +275,12 @@ def simplify(obj):
 
 class JSONSchema(object):
     @classmethod
-    def load(cls, reader, allow_external_refs=False, base_uri=""):
+    def load(cls, reader, allow_external_refs=False, base_uri="", loader=None):
         args = {}
         if not allow_external_refs:
             args['loader'] = custom_jsonref_jsonloader
+        if loader:
+            args['loader'] = loader
         obj = jsonref.load(reader, object_pairs_hook=OrderedDict, base_uri=base_uri, **args)
         return cls.instantiate(None, obj)
 
@@ -255,12 +290,14 @@ def loads(cls, string):
         return cls.instantiate(None, obj)
 
     @classmethod
-    def loadfromfile(cls, filename, allow_external_refs=False):
+    def loadfromfile(cls, filename, allow_external_refs=False, loader=None):
         with io.open(filename, 'rt', encoding='utf-8') as reader:
             args = {}
             if allow_external_refs:
                 args['allow_external_refs'] = True
                 args['base_uri'] = Path(os.path.realpath(filename)).as_uri()
+                if loader:
+                    args['loader'] = loader
             return cls.load(reader, **args)
 
     @classmethod