From c51ebc4b8fd665e6d7e760117a2bc5e555e48078 Mon Sep 17 00:00:00 2001
From: "Pablo R. Mier" <pablo.rodriguez.mier@gmail.com>
Date: Thu, 9 May 2024 14:14:46 +0200
Subject: [PATCH] Improve plotting capabilities

---
 corneto/_graph.py        | 20 ++++++++++++++++-
 corneto/contrib/_util.py | 46 ++++++++++++++++++++++++++++------------
 2 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/corneto/_graph.py b/corneto/_graph.py
index 63e98a1a..f7f3c120 100644
--- a/corneto/_graph.py
+++ b/corneto/_graph.py
@@ -493,7 +493,25 @@ def prune(
         return self.subgraph(reachable)
 
     def plot(self, **kwargs):
-        return self.to_graphviz(**kwargs)
+        Gv = self.to_graphviz(**kwargs)
+        try:
+            return Gv
+        except Exception as e:
+            from corneto._settings import LOGGER
+            from corneto._util import supports_html
+
+            LOGGER.warning(f"SVG+XML rendering failed: {e}.")
+            # Detect if HTML support
+            if supports_html():
+                from corneto.contrib._util import dot_vizjs_html
+
+                class _VizJS:
+                    def _repr_html_(self):
+                        return dot_vizjs_html(Gv)
+
+                return _VizJS()
+            else:
+                raise e
 
     def to_graphviz(self, **kwargs):
         from corneto._plotting import to_graphviz
diff --git a/corneto/contrib/_util.py b/corneto/contrib/_util.py
index 7c6c7d8e..d05b3e8e 100644
--- a/corneto/contrib/_util.py
+++ b/corneto/contrib/_util.py
@@ -1,4 +1,4 @@
-def dot_vizjs(
+def dot_vizjs_html(
     dot_input,
     container_id=None,
     viz_js_url=None,
@@ -8,15 +8,6 @@ def dot_vizjs(
     import base64
     import uuid
 
-    try:
-        from IPython.display import HTML, display
-    except ImportError as e:
-        raise ImportError(
-            "IPython is not installed but is required for displaying the output "
-            "directly in Jupyter notebooks. Please install IPython with "
-            "`pip install ipython`."
-        ) from e
-
     # Generate a random container ID if none is provided
     if container_id is None:
         container_id = f"container-{uuid.uuid4()}"
@@ -40,15 +31,17 @@ def dot_vizjs(
 
     # Base64 encode the DOT content to safely embed it in HTML/JavaScript
     dot_string_base64 = base64.b64encode(dot_string.encode()).decode("utf-8")
-    if vizjs_version:
+    if vizjs_version is not None:
         vizjs_version = f"@{vizjs_version}"
+    else:
+        vizjs_version = ""
     # Setting default URLs if custom URLs are not provided
     if not viz_js_url:
         viz_js_url = f"https://unpkg.com/viz.js{vizjs_version}/viz.js"
     if not full_render_js_url:
         full_render_js_url = f"https://unpkg.com/viz.js{vizjs_version}/full.render.js"
 
-    html_code = f"""
+    return f"""
     <div id='{container_id}'></div> <!-- Container to display the graph -->
     <script>
     function loadScript(url, callback) {{
@@ -81,4 +74,31 @@ def dot_vizjs(
     renderGraph("{dot_string_base64}");
     </script>
     """
-    display(HTML(html_code))
+
+
+def dot_vizjs(
+    dot_input,
+    container_id=None,
+    viz_js_url=None,
+    full_render_js_url=None,
+    vizjs_version=None,
+):
+    html_code = dot_vizjs_html(
+        dot_input,
+        container_id=container_id,
+        viz_js_url=viz_js_url,
+        full_render_js_url=full_render_js_url,
+        vizjs_version=vizjs_version,
+    )
+
+    try:
+        from IPython.display import HTML, display
+
+        display(HTML(html_code))
+
+    except ImportError as e:
+        raise ImportError(
+            "IPython is not installed but is required for displaying the output "
+            "directly in Jupyter notebooks. Please install IPython with "
+            "`pip install ipython`."
+        ) from e