speedup runpython execution for graphviz: cache script results in Sphinx env

Copilot · xadupre · Copilot · commit 77468b7bb292 · 2026-03-03T22:46:20.000Z
Co-authored-by: xadupre &lt;22452781+xadupre@users.noreply.github.com&gt;
diff --git a/_unittests/ut_gdot/test_gdot_extension.py b/_unittests/ut_gdot/test_gdot_extension.py
@@ -137,5 +137,34 @@ def test_gdot4_png(self):
         self.assertIn("png", content)
 
 
+    @ignore_warnings(PendingDeprecationWarning)
+    def test_gdot_script_cache(self):
+        """Test that identical scripts are cached and produce the same output."""
+        script = "print('digraph foo { HbarH -> HbazH; }')".replace("H", '\"')
+        content = f"""
+before
+
+.. gdot::
+    :script:
+
+    {script}
+
+middle
+
+.. gdot::
+    :script:
+
+    {script}
+
+after
+"""
+        content = rst2html(
+            content, writer_name="rst", new_extensions=["sphinx_runpython.gdot"]
+        )
+        # Both gdot directives should produce the same DOT output
+        count = content.count('digraph foo { "bar" -> "baz"; }')
+        self.assertEqual(count, 2, f"Expected the DOT code to appear twice, got {count}")
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/sphinx_runpython/gdot/sphinx_gdot_extension.py b/sphinx_runpython/gdot/sphinx_gdot_extension.py
@@ -1,3 +1,4 @@
+import hashlib
 import os
 import logging
 from docutils import nodes
@@ -147,7 +148,23 @@ def run(self):
         # executes script if any
         content = "\n".join(self.content)
         if script or script == "":
-            stdout, stderr, _ = run_python_script(content, process=process)
+            env = info.get("env")
+            cache_key = hashlib.sha256(
+                f"{content}:{process}".encode("utf-8")
+            ).hexdigest()
+            if env is not None:
+                if not hasattr(env, "gdot_script_cache"):
+                    env.gdot_script_cache = {}
+                cached = env.gdot_script_cache.get(cache_key, None)
+            else:
+                cached = None
+
+            if cached is not None:
+                stdout, stderr = cached
+            else:
+                stdout, stderr, _ = run_python_script(content, process=process)
+                if env is not None:
+                    env.gdot_script_cache[cache_key] = (stdout, stderr)
 
             if stderr:
                 out = [
