From 3666fd458a14a6466c07802eda56cabba97c14c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Davide=20Sandon=C3=A0?= <sandona.davide@gmail.com>
Date: Mon, 24 Oct 2022 23:27:15 +0200
Subject: [PATCH] Added capability to save png or other static format

As I'm looking to use this Sphinx extension on my module, considering the
high number of plots (many of which are 3D plots that requires a lot of
bandwidth and really slows down the rendering of the page), I figured
that it would be nice to visualize png plots.

This commits modifies `save_plotly_figure` and add a couple of tests to
verify its correct behaviour.

Users can set the following on `conf.py` (for example):

plotly_formats = ["png", "html", "pdf"]

This will visualize png by default, also adding links to a html and pdf
versions.
---
 sphinx_plotly_directive/utils.py | 11 ++++++++---
 tests/test_utils.py              |  8 ++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/sphinx_plotly_directive/utils.py b/sphinx_plotly_directive/utils.py
index 9d6d8ec..ec7daeb 100644
--- a/sphinx_plotly_directive/utils.py
+++ b/sphinx_plotly_directive/utils.py
@@ -30,9 +30,14 @@ def save_plotly_figure(fig, path):
     >>> path = tempfile.NamedTemporaryFile(suffix=".html").name
     >>> save_plotly_figure(fig, path)
     """
-    fig_html = plotly.offline.plot(fig, output_type="div", include_plotlyjs="cdn", auto_open=False)
-    with open(path, "w") as f:
-        f.write(fig_html)
+    ext = path.split(".")[-1]
+    if ext in ["htm", "html"]:
+        fig_html = plotly.offline.plot(fig, output_type="div",
+            include_plotlyjs="cdn", auto_open=False)
+        with open(path, "w") as f:
+            f.write(fig_html)
+    else:
+        fig.write_image(path)
 
 
 def assign_last_line_into_variable(code, variable_name):
diff --git a/tests/test_utils.py b/tests/test_utils.py
index d25a54f..f9c2b92 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -19,6 +19,14 @@ def tset_save_plotly_figure(tmpdir):
     save_plotly_figure(fig, out_path)
     assert os.path.exists(out_path)
 
+    out_path = os.path.join(tmpdir.strpath, "fig.png")
+    save_plotly_figure(fig, out_path)
+    assert os.path.exists(out_path)
+
+    out_path = os.path.join(tmpdir.strpath, "fig.pdf")
+    save_plotly_figure(fig, out_path)
+    assert os.path.exists(out_path)
+
 
 def test_assign_last_line_into_variable():
     code = """