improve versatility

sdpython · sdpython · Jan 8, 2024 · Jan 8, 2024 · Jan 8, 2024 · Jan 8, 2024
commit 8436ea447246d57e782866e4f099d93313db263e
diff --git a/.gitignore b/.gitignore
@@ -14,6 +14,7 @@ build/*
 *egg-info/*
 onnxruntime_profile*
 prof
+test*.png
 _doc/sg_execution_times.rst
 _doc/auto_examples/*
 _doc/examples/_cache/*

diff --git a/_unittests/ut_plotting/test_graphviz.py b/_unittests/ut_plotting/test_graphviz.py
@@ -25,6 +25,12 @@ def test_draw_graph_graphviz(self):
         draw_graph_graphviz(dot, image=fout)
         self.assertExists(os.path.exists(fout))
 
+    def test_draw_graph_graphviz_proto(self):
+        fout = "test_draw_graph_graphviz_proto.png"
+        dot = self._get_graph()
+        draw_graph_graphviz(dot, image=fout)
+        self.assertExists(os.path.exists(fout))
+
     def test_plot_dot(self):
         dot = to_dot(self._get_graph())
         ax = plot_dot(dot)

diff --git a/onnx_array_api/plotting/graphviz_helper.py b/onnx_array_api/plotting/graphviz_helper.py
@@ -2,8 +2,9 @@
 import subprocess
 import sys
 import tempfile
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 import numpy as np
+from onnx import ModelProto
 
 
 def _find_in_PATH(prog: str) -> Optional[str]:
@@ -139,24 +140,29 @@ def _run_graphviz(filename: str, image: str, engine: str = "dot") -> str:
 
 
 def draw_graph_graphviz(
-    dot: str,
+    dot: Union[str, ModelProto],
     image: str,
     engine: str = "dot",
 ) -> str:
     """
     Draws a graph using :epkg:`Graphviz`.
 
-    :param dot: dot graph
+    :param dot: dot graph or ModelProto
     :param image: output image, None, just returns the output
     :param engine: *dot* or *neato*
     :return: :epkg:`Graphviz` output or
         the dot text if *image* is None
 
     The function creates a temporary file to store the dot file if *image* is not None.
     """
+    if isinstance(dot, ModelProto):
+        from .dot_plot import to_dot
+
+        sdot = to_dot(dot)
+    else:
+        sdot = dot
     with tempfile.NamedTemporaryFile(delete=False) as fp:
-        fp.write(dot.encode("utf-8"))
-        fp.seek(0)
+        fp.write(sdot.encode("utf-8"))
         fp.close()
 
         filename = fp.name
@@ -172,15 +178,15 @@ def draw_graph_graphviz(
 
 
 def plot_dot(
-    dot: str,
+    dot: Union[str, ModelProto],
     ax: Optional["matplotlib.axis.Axis"] = None,  # noqa: F821
     engine: str = "dot",
     figsize: Optional[Tuple[int, int]] = None,
 ) -> "matplotlib.axis.Axis":  # noqa: F821
     """
     Draws a dot graph into a matplotlib graph.
 
-    :param dot: dot graph
+    :param dot: dot graph or ModelProto
     :param image: output image, None, just returns the output
     :param engine: *dot* or *neato*
     :param figsize: figsize of ax is None