postprocessing with LinkML SchemaView

main.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import subprocess
 
 from linkml.generators.jsonschemagen import JsonSchemaGenerator
@@ -7,7 +8,9 @@
 from linkml.generators.jsonldcontextgen import ContextGenerator
 from linkml.generators.docgen import DocGenerator
 from linkml.generators.linkmlgen import LinkmlGenerator
+from linkml_runtime.utils.schemaview import SchemaView
 from pathlib import Path
+from pyld import jsonld
 
 RESOURCES_PATH = Path('resources')
 GENS_PATH = RESOURCES_PATH / 'gens'
@@ -25,28 +28,44 @@ def generate_docs():
     doc_generator = DocGenerator(YAML_SCHEMA_PATH, mergeimports=False)
     doc_generator.serialize(directory=str(DOCDIR))
 
+
+def post_process_jsonldcontext(schema_view: SchemaView, serialized_schema: str) -> str:
+    serialized_schema_json = json.loads(serialized_schema)
+    for slot in schema_view.all_slots().values():
+        if slot.slot_uri is not None and 'InLanguage' in slot.slot_uri and isinstance(serialized_schema_json['@context'][slot.name], dict):
+            serialized_schema_json['@context'][slot.name]['@container'] = '@language'
+        if slot.multivalued and str(slot.range) == 'Any' and isinstance(serialized_schema_json['@context'][slot.name], dict):
+            if '@type' in serialized_schema_json['@context'][slot.name].keys():
+                del serialized_schema_json['@context'][slot.name]['@type']
+            serialized_schema_json['@context'][slot.name]['@container'] = '@set'
+    return json.dumps(serialized_schema_json, indent=3)
+
+
 def main(generate_docs_flag, serve_docs_flag):
     if not YAML_SCHEMA_PATH.exists():
         print(f"LinkML schema file does not exist: {YAML_SCHEMA_PATH}")
         return
+    linkml_schema_view = SchemaView(YAML_SCHEMA_PATH, merge_imports=True)
+    # TODO: add pre processing for LinkML if needed
     for generator in GENERATORS:
         output_dir = GENS_PATH / generator
         output_dir.mkdir(parents=True, exist_ok=True)
         if generator == 'jsonschema':
             # json_schema_generator = JsonSchemaGenerator(yaml_content, top_class="Thing")
-            json_schema_generator = JsonSchemaGenerator(YAML_SCHEMA_PATH, mergeimports=True)
+            json_schema_generator = JsonSchemaGenerator(linkml_schema_view.schema, mergeimports=True)
             (output_dir / 'jsonschema.json').write_text(json_schema_generator.serialize())
         elif generator == 'shacl':
-            shacl_generator = ShaclGenerator(YAML_SCHEMA_PATH, mergeimports=False, closed=True, suffix='Shape')
+            shacl_generator = ShaclGenerator(linkml_schema_view.schema, mergeimports=False, closed=True, suffix='Shape')
             (output_dir / 'shapes.shacl.ttl').write_text(shacl_generator.serialize())
         elif generator == 'owl':
-            owl_generator = OwlSchemaGenerator(YAML_SCHEMA_PATH,)
+            owl_generator = OwlSchemaGenerator(linkml_schema_view.schema,)
             (output_dir / 'ontology.owl.ttl').write_text(owl_generator.serialize())
         elif generator == 'jsonldcontext':
-            context_generator = ContextGenerator(YAML_SCHEMA_PATH, mergeimports=False)
-            (output_dir / 'context.jsonld').write_text(context_generator.serialize())
+            context_generator = ContextGenerator(linkml_schema_view.schema, mergeimports=True)
+            (output_dir / 'context.jsonld').write_text(post_process_jsonldcontext(linkml_schema_view,
+                                                                                  context_generator.serialize()))
         elif generator == 'linkml':
-            linkml_generator = LinkmlGenerator(YAML_SCHEMA_PATH, mergeimports=True, format='yaml', output='linkml.yaml')
+            linkml_generator = LinkmlGenerator(linkml_schema_view.schema, mergeimports=True, format='yaml', output='linkml.yaml')
             (output_dir / 'linkml.yaml').write_text(linkml_generator.serialize())
         else:
             print(f"Unknown generator: {generator}")
@@ -69,3 +88,4 @@ def main(generate_docs_flag, serve_docs_flag):
                         help='Boolean for serving the generated documentation.')
     args = parser.parse_args()
     main(args.local_docs, args.serve_docs)
+