Generate diagram for Session base class

.github/workflows/docs.yml

@@ -9,6 +9,8 @@ jobs:
   pages:
     runs-on: ubuntu-20.04
     steps:
+    - name: Setup Graphviz
+      uses: ts-graphviz/[email protected]
     - id: deployment
       uses: sphinx-notes/pages@v3
       with:

docs/api.base/session.rst

@@ -1,5 +1,8 @@
 AindBehaviorSessionModel
--------------
+--------------------------
+
+.. image:: ../_static/AindBehaviorSessionModel.svg
+   :target: ../_static/AindBehaviorSessionModel.svg
 
 .. autopydantic_model:: aind_behavior_services.session.AindBehaviorSessionModel
    :members:

docs/conf.py

@@ -9,19 +9,20 @@
 
 import os
 import sys
-
-sys.path.insert(0, os.path.abspath("../src/DataSchemas"))
+import erdantic as erd
+from pydantic import BaseModel
 import glob
 
-from aind_behavior_services import __version__
+sys.path.insert(0, os.path.abspath("../src/DataSchemas"))
+import aind_behavior_services
 
 SOURCE_ROOT = "https://github.com/AllenNeuralDynamics/Aind.Behavior.Services/tree/main/src/DataSchemas/"
 
 
 project = "AIND Behavior Services"
 copyright = "2024, Allen Institute for Neural Dynamics"
 author = "Bruno Cruz"
-release = __version__
+release = aind_behavior_services.__version__
 
 
 # -- Generate jsons --------------------------------------------------------------
@@ -40,7 +41,7 @@
 
 leaf_template = """
 {json_file_name}
--------------
+----------------------------------------------------
 
 `Download Schema <https://raw.githubusercontent.com/AllenNeuralDynamics/Aind.Behavior.Services/main/src/DataSchemas/schemas/{json_file_name}.json>`_
 
@@ -111,3 +112,14 @@ def linkcode_resolve(domain, info):
         return None
     filename = info["module"].replace(".", "/")
     return f"{SOURCE_ROOT}/{filename}.py"
+
+
+# -- Class diagram generation
+
+def export_model_diagram(model: BaseModel, root: str = "_static") -> None:
+    diagram = erd.create(model)
+    diagram.draw(f"{root}/{model.__name__}.svg")
+
+
+_diagram_root = "_static"
+export_model_diagram(aind_behavior_services.session.AindBehaviorSessionModel, _diagram_root)

pyproject.toml

@@ -41,7 +41,7 @@ docs = [
     'Sphinx<7.3',
     'furo',
     'sphinx-jinja',
-    'autodoc_pydantic',
+    'autodoc_pydantic[erdantic]',
     'myst_parser',
     'sphinx-jsonschema',
     'sphinx-copybutton'