adds option to return html, svg, png as bytes

cloudbuild.yaml

@@ -1,4 +1,4 @@
 steps:
 - name: 'gcr.io/cloud-builders/docker'
-  args: ['build', '-t', 'gcr.io/$PROJECT_ID/draw-base-image:0.19', '.'] 
-images: ['gcr.io/$PROJECT_ID/draw-base-image:0.19']
+  args: ['build', '-t', 'gcr.io/$PROJECT_ID/draw-base-image:0.20', '.'] 
+images: ['gcr.io/$PROJECT_ID/draw-base-image:0.20']

requirements.txt

@@ -20,4 +20,5 @@ sphinxcontrib-blockdiag
 sphinx_panels
 sphinx_click
 xmltodict
-pytest
+pytest
+kaleido

seismic_graph/util/misc.py

@@ -135,26 +135,52 @@ def _generate_legend(self, fit, m, prefix, suffix):
 
 
 
-
-def save_plot(func,  *args, **kwargs):
-    """Default arguments for saving a plotly figure.
+def save_plot(func):
+    """Arguments for handling and saving a Plotly figure.
 
     Args:
-        to_html (str, optional): File name to save the figure as a html.
-        to_png (str, optional): File name to save the figure as a png.
+        to_html (str, optional): File name or 'return' to save the figure as HTML or return HTML as bytes.
+        to_png (str, optional): File name or 'return' to save the figure as PNG or return PNG as bytes.
+        to_svg (str, optional): File name or 'return' to save the figure as SVG or return SVG as bytes.
     """
-
     def wrapper(*args, **kwargs):
         out = func(*args, **kwargs)
-        to_html, to_png = kwargs.pop('to_html', None), kwargs.pop('to_png', None)
+        fig = out['fig']
+
+        # Handle HTML output
+        to_html = kwargs.pop('to_html', None)
         if to_html:
-            out['fig'].write_html(to_html, include_plotlyjs='cdn')
+            if to_html == 'return':
+                html_bytes = fig.to_html(include_plotlyjs='cdn', full_html=False).encode('utf-8')
+                out['html_bytes'] = html_bytes
+            else:
+                fig.write_html(to_html, include_plotlyjs='cdn')
+
+        # Handle PNG output
+        to_png = kwargs.pop('to_png', None)
         if to_png:
-            out['fig'].write_image(to_png)
+            if to_png == 'return':
+                png_bytes = fig.to_image(format="png")
+                out['png_bytes'] = png_bytes
+            else:
+                fig.write_image(to_png)
+
+        # Handle SVG output
+        to_svg = kwargs.pop('to_svg', None)
+        if to_svg:
+            if to_svg == 'return':
+                svg_bytes = fig.to_image(format="svg")
+                out['svg_bytes'] = svg_bytes
+            else:
+                fig.write_image(to_svg, format='svg')
+
         return out
-    wrapper.__name__=func.__name__
-    wrapper.__doc__=func.__doc__
-    return wrapper   
+
+    # Preserve function metadata
+    wrapper.__name__ = func.__name__
+    wrapper.__doc__ = func.__doc__
+    return wrapper
+
 
 def extract_args(func):
     """Extract the arguments of a function.