sphinxcontrib/jsonschema.py: Add urn handling

sphinxcontrib/jsonschema.py

@@ -13,6 +13,8 @@
 from docutils import nodes
 from docutils.parsers.rst import directives, Directive
 from pathlib import Path
+from urllib import parse as urlparse
+from functools import partial
 
 import json
 from collections import OrderedDict
@@ -38,6 +40,21 @@ def custom_jsonref_jsonloader(uri, **kwargs):
     return {}
 
 
+def urnjsonloader(load_path, uri, **kwargs):
+    """
+    Provides a callable which handles a URN if provided, else calls default JSON loader.
+    A partial needs to be created to set the load path, before passing to jsonref.
+    """
+    scheme = urlparse.urlsplit(uri).scheme
+
+    if scheme == "urn":
+        with open(f'{load_path}/{uri.split("urn:")[-1]}.json') as content:
+            result = json.loads(content.read(), **kwargs)
+        return result
+    else:
+        return jsonref.jsonloader(uri, **kwargs)
+
+
 class JSONSchemaDirective(Directive):
     has_content = True
     required_arguments = 1
@@ -49,6 +66,7 @@ class JSONSchemaDirective(Directive):
         'addtargets': directives.flag,
         'externallinks': directives.unchanged,
         'allowexternalrefs': directives.flag,
+        'allowurnrefs': directives.flag,
     }
     # Add a rollup option here
 
@@ -87,11 +105,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=(partial(urnjsonloader, 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=(partial(urnjsonloader, 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 +267,11 @@ 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 +281,13 @@ 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