wip: clean up + rst inheritance

examples/Allele.json

@@ -55,23 +55,17 @@
           "$ref": "/ga4gh/schema/gks-common/1.x/core-im/json/IRI"
         },
         {
-          "$ref": "/ga4gh/schema/vrs/2.x/json/SequenceLocation"
+          "$ref": "/ga4gh/schema/vrs/2.x/json/Location"
         }
       ]
     },
     "state": {
-      "description": "An expression of the sequence state",
-      "oneOf": [
-        {
-          "$ref": "/ga4gh/schema/vrs/2.x/json/LiteralSequenceExpression"
-        },
+      "allOf": [
         {
-          "$ref": "/ga4gh/schema/vrs/2.x/json/ReferenceLengthExpression"
-        },
-        {
-          "$ref": "/ga4gh/schema/vrs/2.x/json/LengthExpression"
+          "$ref": "/ga4gh/schema/vrs/2.x/json/SequenceExpression"
         }
-      ]
+      ],
+      "description": "An expression of the sequence state"
     }
   },
   "required": [

examples/Allele.rst

@@ -4,6 +4,7 @@ The state of a molecule at a :ref:`Location`.
 
 **Information Model**
 
+Some Allele attributes are inherited from :ref:`Variation`.
 
 .. list-table::
    :class: clean-wrap
@@ -48,10 +49,10 @@ The state of a molecule at a :ref:`Location`.
      - 0..m
      - None
    * - location
-     - :ref:`IRI` | :ref:`SequenceLocation`
+     - :ref:`IRI` | :ref:`Location`
      - 0..1
      - The location of the Allele
    * - state
-     - :ref:`LiteralSequenceExpression` | :ref:`ReferenceLengthExpression` | :ref:`LengthExpression`
+     - :ref:`SequenceExpression`
      - 0..1
      - An expression of the sequence state

examples/pydantic_to_json_schema.py

@@ -1,8 +1,7 @@
 from enum import Enum
-from typing import List, Literal, Optional
 import json
 
-from pydantic import Field, BaseModel, RootModel
+from pydantic import BaseModel, RootModel
 from pydantic.json_schema import GenerateJsonSchema
 
 from ga4gh.core import entity_models, domain_models
@@ -17,9 +16,9 @@
 }
 
 
-def create_model_module_map(*modules) -> dict[str, str]:
-    """Creates a mapping from model names to their modules."""
-    model_module_map = {}
+def map_model_to_ref(*modules) -> dict[str, str]:
+    """Creates a mapping from model names to their JSON schema references."""
+    model_to_module = {}
     for module in modules:
         for attr_name in dir(module):
             model = getattr(module, attr_name)
@@ -28,33 +27,11 @@ def create_model_module_map(*modules) -> dict[str, str]:
                 and issubclass(model, (BaseModel, RootModel, Enum))
                 and model.__module__ == module.__name__
             ):
-                model_module_map[attr_name] = MODULE_TO_REF[model.__module__]
-    return model_module_map
+                model_to_module[attr_name] = MODULE_TO_REF[model.__module__]
+    return model_to_module
 
 
-MODEL_REF_MAP = create_model_module_map(domain_models, entity_models, models)
-
-
-class Allele(models.Allele, extra="forbid"):
-    """The state of a molecule at a :ref:`Location`."""
-
-    class Config:
-        @staticmethod
-        def json_schema_extra(cls):
-            for prop in {"location", "state"}:
-                cls["properties"][prop]["oneOf"] = cls["properties"][prop]["anyOf"]
-                del cls["properties"][prop]["anyOf"]
-
-    expressions: Optional[List[models.Expression]] = Field(None, ordered=False)
-    maturity: Literal["draft"] = Field("draft", frozen=True)
-    alternativeLabels: Optional[List[str]] = Field(
-        None, description="Alternative name(s) for the Entity.", ordered=False
-    )
-    extensions: Optional[List[entity_models.Extension]] = Field(
-        None,
-        description="A list of extensions to the Entity, that allow for capture of information not directly supported by elements defined in the model.",
-        ordered=False,
-    )
+MODEL_TO_REF = map_model_to_ref(domain_models, entity_models, models)
 
 
 class GksGenerateJsonSchema(GenerateJsonSchema):
@@ -81,7 +58,7 @@ def traverse_and_modify(self, schema):
 
             if "$ref" in schema:
                 class_name = schema["$ref"].split("/")[-1]
-                schema["$ref"] = f"{MODEL_REF_MAP[class_name]}/{class_name}"
+                schema["$ref"] = f"{MODEL_TO_REF[class_name]}/{class_name}"
 
             if "description" in schema and isinstance(schema["description"], str):
                 schema["description"] = scrub_rst_markup(schema["description"])
@@ -134,7 +111,7 @@ def generate(self, schema, mode="validation"):
 if __name__ == "__main__":
     with open("examples/Allele.json", "w") as wf:
         json.dump(
-            Allele.model_json_schema(schema_generator=GksGenerateJsonSchema),
+            models.Allele.model_json_schema(schema_generator=GksGenerateJsonSchema),
             wf,
             indent=2,
         )

examples/pydantic_to_rst.py

@@ -5,8 +5,9 @@
 from typing import Annotated, Any, List
 from typing import get_args, get_origin
 
-from utils import EXCLUDE_PROPS
-from pydantic_to_json_schema import Allele
+from utils import EXCLUDE_PROPS, INSTANCE_TO_ABC
+
+from ga4gh.vrs import models
 
 
 PYTHON_TO_JSON_TYPES = {
@@ -61,8 +62,11 @@ def get_limits(
 
 
 def generate(model: BaseModel) -> str:
-    # TODO: Inheritance
-    inheritance = ""
+    model_name = model.__name__
+    if model_name in INSTANCE_TO_ABC:
+        inheritance = f"Some {model_name} attributes are inherited from :ref:`{INSTANCE_TO_ABC[model_name]}`.\n"
+    else:
+        inheritance = ""
 
     rst_data = [
         "**Computational Definition**",
@@ -95,11 +99,15 @@ def generate(model: BaseModel) -> str:
             field_is_list = False
             field_type = f":ref:`{field_info.annotation.__name__}`"
         else:
-            field_annotation = tuple(
-                anno
-                for anno in get_args(field_info.annotation)
-                if anno is not type(None)
-            )
+            field_anno_args = get_args(field_info.annotation)
+            if field_anno_args:
+                field_annotation = tuple(
+                    anno
+                    for anno in get_args(field_info.annotation)
+                    if anno is not type(None)
+                )
+            else:
+                field_annotation = [field_info.annotation]
 
             field_is_list = get_origin(field_annotation[0]) in {list, List}
 
@@ -131,4 +139,4 @@ def generate(model: BaseModel) -> str:
 
 if __name__ == "__main__":
     with open("examples/Allele.rst", "w") as wf:
-        wf.write(generate(Allele))
+        wf.write(generate(models.Allele))

examples/utils.py

@@ -1,4 +1,10 @@
 import re
+from typing import get_args
+
+from pydantic import RootModel
+
+from ga4gh.core import entity_models, domain_models
+from ga4gh.vrs import models
 
 
 REF_RE = re.compile(r":ref:`(.*?)(\s?<.*>)?`")
@@ -11,3 +17,57 @@ def scrub_rst_markup(string):
     string = LINK_RE.sub(r"[\g<1>](\g<2>)", string)
     string = string.replace("\n", " ")
     return string
+
+
+def map_abc_to_instances(*modules) -> dict[str, str]:
+    """Creates a mapping from ABC model names to their instance model names."""
+    abc_to_instances = {}
+    excluded_types = {
+        "list",
+        "str",
+        "Optional",
+        "int",
+        "float",
+        "dict",
+        "bool",
+        "set",
+        "tuple",
+    }
+
+    for module in modules:
+        for attr_name in dir(module):
+            model = getattr(module, attr_name)
+            if (
+                isinstance(model, type)
+                and issubclass(model, RootModel)
+                and model.__module__ == module.__name__
+            ):
+                root_anno = model.model_fields["root"].annotation
+                root_annos = get_args(root_anno) or (root_anno,)
+                root_anno_cls_names = [
+                    cls.__name__
+                    for cls in root_annos
+                    if cls.__name__ not in excluded_types
+                ]
+                if root_anno_cls_names:
+                    abc_to_instances[model.__name__] = root_anno_cls_names
+
+    return abc_to_instances
+
+
+def get_abc(key) -> str:
+    """Get original ABC class name
+
+    :param key: Class name
+    :return: Original ABC class name
+    """
+    while key in INSTANCE_TO_ABC:
+        key = INSTANCE_TO_ABC[key]
+    return key
+
+
+ABC_TO_INSTANCES = map_abc_to_instances(models, entity_models, domain_models)
+INSTANCE_TO_ABC = {
+    value: key for key, value_list in ABC_TO_INSTANCES.items() for value in value_list
+}
+INSTANCE_TO_ABC = {key: get_abc(key) for key in INSTANCE_TO_ABC}