From 78eb0e4992cfefd04d8739f58b8f6902de589860 Mon Sep 17 00:00:00 2001
From: glopesdev <glopesdev@users.noreply.github.com>
Date: Wed, 18 Dec 2024 13:51:10 +0000
Subject: [PATCH] Ensure workflow files are hosted in fixed path

---
 src/_static/js/workflow.js |  7 ++--
 src/conf.py                |  2 ++
 src/convertworkflow.py     | 69 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 3 deletions(-)
 create mode 100644 src/convertworkflow.py

diff --git a/src/_static/js/workflow.js b/src/_static/js/workflow.js
index d0549186..83c15817 100644
--- a/src/_static/js/workflow.js
+++ b/src/_static/js/workflow.js
@@ -45,9 +45,10 @@ class WorkflowContainer {
         element.classList.add("highlight-bonsai");
         const img = element.querySelector("img");
 
-        // assumes workflow references with $ are hosted in _static/
-        const workflowPath = img.src.replace("$", "_static/");
-        img.src = workflowPath.replace(/\.[^.]+$/, ".svg");
+        // assumes all workflow references are hosted in _workflows/
+        const workflowPath = img.src
+            .replace("_images/", "_workflows/")
+            .replace(/\.[^.]+$/, ".bonsai");
 
         const wrap = WorkflowContainer.createWorkflowCell(img, workflowPath);
         const parent = element.parentElement;
diff --git a/src/conf.py b/src/conf.py
index df8b765e..cc146635 100644
--- a/src/conf.py
+++ b/src/conf.py
@@ -17,6 +17,7 @@
 
 sys.path.extend(
     [
+        os.path.abspath("."),
         os.path.abspath("../aeon_mecha/"),
         os.path.abspath("../aeon_analysis/aeon_analysis"),
         os.path.abspath("../aeon_acquisition"),
@@ -68,6 +69,7 @@ def get_current_release_tag():
     "myst_nb",
     "sphinx_design",
     "sphinx_copybutton",
+    "convertworkflow"
 ]
 
 # Configure the myst parser to enable cool markdown features
diff --git a/src/convertworkflow.py b/src/convertworkflow.py
new file mode 100644
index 00000000..8cc1d5af
--- /dev/null
+++ b/src/convertworkflow.py
@@ -0,0 +1,69 @@
+import os
+import sphinx
+from pathlib import Path
+from docutils import nodes
+
+from sphinx.locale import __
+from sphinx.util import logging
+from sphinx.util.osutil import ensuredir, copyfile
+
+from sphinx.application import Sphinx
+from sphinx.util.typing import ExtensionMetadata
+
+from sphinx.transforms.post_transforms.images import BaseImageConverter
+
+logger = logging.getLogger(__name__)
+
+
+class WorkflowImageConverter(BaseImageConverter):
+    """An image converter to resolve workflow image references.
+    
+    The converter will handle any image reference targeting a .bonsai
+    workflow file and look for the corresponding .svg file.
+
+    If the builder output is html it will also copy the source file
+    into the _workflows folder so it can be fetched by copy button scripts.
+    """
+    default_priority = 300
+
+    def match(self, node: nodes.image) -> bool:
+        _, ext = os.path.splitext(node['uri'])
+        return '://' not in node['uri'] and ext == '.bonsai'
+    
+    def handle(self, node: nodes.image) -> None:
+        try:
+            srcpath = Path(node['uri'])
+            abs_srcpath = self.app.srcdir / srcpath
+            abs_imgpath = abs_srcpath.with_suffix('.svg')
+            if not abs_imgpath.exists():
+                logger.warning(__('Could not find workflow image: %s [%s]'), node['uri'], abs_imgpath)
+                return
+
+            # copy workflow image to _images folder
+            destpath = os.path.join(self.imagedir, abs_imgpath.name)
+            ensuredir(self.imagedir)
+            copyfile(abs_imgpath, destpath)
+
+            # resolve image cross-references
+            if '*' in node['candidates']:
+                node['candidates']['*'] = destpath
+            node['uri'] = destpath
+            self.env.original_image_uri[destpath] = srcpath
+            self.env.images.add_file(self.env.docname, destpath)
+
+            # copy workflow file to _workflows folder when output is html
+            if self.app.builder.format == 'html':
+                abs_workflowdir = self.app.builder.outdir / '_workflows'
+                abs_workflowpath = abs_workflowdir / abs_srcpath.name
+                ensuredir(abs_workflowdir)
+                copyfile(abs_srcpath, abs_workflowpath)
+        except Exception as exc:
+            logger.warning(__('Could not fetch workflow image: %s [%s]'), node['uri'], exc)
+
+def setup(app: Sphinx) -> ExtensionMetadata:
+    app.add_post_transform(WorkflowImageConverter)
+    return {
+        'version': sphinx.__display_version__,
+        'parallel_read_safe': True,
+        'parallel_write_safe': True,
+    }
\ No newline at end of file