Further manual fixes to get things to work

gyorilab · Sep 7, 2024 · 0bd9ef6 · 0bd9ef6
1 parent 7491d52
commit 0bd9ef6
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 25 deletions.
diff --git a/mira/metamodel/schema.py b/mira/metamodel/schema.py
@@ -6,7 +6,8 @@
 from pathlib import Path
 
 import pydantic
-from pydantic import BaseModel
+from pydantic import BaseModel, create_model
+from pydantic.json_schema import model_json_schema
 
 from . import Concept, Template, TemplateModel
 
@@ -15,29 +16,31 @@
 
 
 def get_json_schema():
-    """Get the JSON schema for MIRA.
-
-    Returns
-    -------
-    : JSON
-        The JSON schema for MIRA.
-    """
+    """Get the JSON schema for MIRA."""
     rv = {
         "$schema": "http://json-schema.org/draft-07/schema#",
         "$id": "https://raw.githubusercontent.com/indralab/mira/main/mira/metamodel/schema.json",
+        "title": "MIRA Metamodel Template Schema",
+        "description": "MIRA metamodel templates give a high-level abstraction of modeling appropriate for many domains.",
     }
-    rv.update(
-        pydantic.schema.schema(
-            [
-                Concept,
-                Template,
-                *Template.__subclasses__(),
-                TemplateModel,
-            ],
-            title="MIRA Metamodel Template Schema",
-            description="MIRA metamodel templates give a high-level abstraction of modeling appropriate for many domains.",
-        )
+
+    models = [Concept, Template, *Template.__subclasses__(), TemplateModel]
+
+    CombinedModel = create_model(
+        "CombinedModel",
+        **{model.__name__: (model, ...) for model in models}
     )
+
+    schema = model_json_schema(
+        CombinedModel,
+        mode='validation',
+    )
+
+    # Remove the top-level 'title' and 'description' from the generated schema
+    schema.pop('title', None)
+    schema.pop('description', None)
+
+    rv.update(schema)
     return rv
 
 

diff --git a/mira/metamodel/template_model.py b/mira/metamodel/template_model.py
@@ -353,13 +353,13 @@ class TemplateModel(BaseModel):
     )
 
     annotations: Optional[Annotations] = Field(
-        default_factory=None,
+        default=None,
         description="A structure containing model-level annotations. "
         "Note that all annotations are optional.",
     )
 
     time: Optional[Time] = Field(
-        default_factory=None,
+        default=None,
         description="A structure containing time-related annotations. "
         "Note that all annotations are optional.",
     )

diff --git a/mira/metamodel/utils.py b/mira/metamodel/utils.py
@@ -50,10 +50,10 @@ def __get_validators__(cls):
     def __get_pydantic_core_schema__(
         cls, source_type: Any, handler: GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
-        return core_schema.union_schema([
+        return handler.resolve_ref_schema(core_schema.union_schema([
             core_schema.is_instance_schema(cls),
             core_schema.no_info_plain_validator_function(cls.validate)
-        ])
+        ]))
 
 
     @classmethod
@@ -67,8 +67,11 @@ def validate(cls, v):
         return cls(v)
 
     @classmethod
-    def __get_pydantic_json_schema__(cls, field_schema):
-        field_schema.update(type="string", example="2*x")
+    def __get_pydantic_json_schema__(cls, core_schema, handler):
+        json_schema = handler(core_schema)
+        json_schema.update(type="string", format="sympy-expr")
+        return json_schema
+        #field_schema.update(type="string", example="2*x")
 
     def __str__(self):
         return super().__str__()[len(self.__class__.__name__)+1:-1]