Add parameter show to plot_onnx (#213)

sdpython · Feb 11, 2021 · 991f12c · 991f12c
1 parent 20bfc14
commit 991f12c
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mlprodict/plotting/plotting_onnx.py b/mlprodict/plotting/plotting_onnx.py
@@ -6,7 +6,8 @@
 from ..onnxrt import OnnxInference
 
 
-def plot_onnx(onx, ax=None, dpi=300, temp_dot=None, temp_img=None):
+def plot_onnx(onx, ax=None, dpi=300, temp_dot=None, temp_img=None,
+              show=False):
     """
     Plots an ONNX graph into graph.
 
@@ -17,6 +18,7 @@ def plot_onnx(onx, ax=None, dpi=300, temp_dot=None, temp_img=None):
         if None, a file is created and removed
     :param temp_img: temporary image,
         if None, a file is created and removed
+    :param show: calls `plt.show()`
     :return: axes
     """
     # delayed import
@@ -30,4 +32,6 @@ def plot_onnx(onx, ax=None, dpi=300, temp_dot=None, temp_img=None):
         oinf = onx
     dot = oinf.to_dot()
     plot_graphviz(dot, dpi=dpi, ax=ax, temp_dot=temp_dot, temp_img=temp_img)
+    if show:
+        plt.show()  # pragma: no cover
     return ax