diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index af2c1ee4..6b8123c8 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -9,6 +9,8 @@ jobs:
pages:
runs-on: ubuntu-20.04
steps:
+ - name: Setup Graphviz
+ uses: ts-graphviz/setup-graphviz@v2.0.2
- id: deployment
uses: sphinx-notes/pages@v3
with:
diff --git a/docs/_static/AindBehaviorSessionModel.svg b/docs/_static/AindBehaviorSessionModel.svg
new file mode 100644
index 00000000..f6879888
--- /dev/null
+++ b/docs/_static/AindBehaviorSessionModel.svg
@@ -0,0 +1,71 @@
+
+
+
+
+
diff --git a/docs/api.base/session.rst b/docs/api.base/session.rst
index c989c793..2335b6a3 100644
--- a/docs/api.base/session.rst
+++ b/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:
diff --git a/docs/conf.py b/docs/conf.py
index 20fa4541..7e59209c 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -9,11 +9,12 @@
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/"
@@ -21,7 +22,7 @@
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 `_
@@ -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)
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index b258475c..761e5774 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -41,7 +41,7 @@ docs = [
'Sphinx<7.3',
'furo',
'sphinx-jinja',
- 'autodoc_pydantic',
+ 'autodoc_pydantic[erdantic]',
'myst_parser',
'sphinx-jsonschema',
'sphinx-copybutton'