Merge pull request #17 from transferwise/feature/return_fig

added return_fig for streamlit
transferwise · Aug 2, 2023 · 22eb6a3 · 22eb6a3
2 parents bea707f + da2be48
commit 22eb6a3
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 30 deletions.
diff --git a/notebooks/Finding interesting segments.ipynb b/notebooks/Finding interesting segments.ipynb
diff --git a/wise_pizza/explain.py b/wise_pizza/explain.py
@@ -304,8 +304,8 @@ def explain_levels(
         s["total"] += average * s["seg_size"]
     # print(average)
     sf.reg.intercept_ = average
-    sf.plot = lambda plot_is_static=False, width=2000, height=500: plot_segments(
-        sf, plot_is_static=plot_is_static, width=width, height=height
+    sf.plot = lambda plot_is_static=False, width=2000, height=500, return_fig=False: plot_segments(
+        sf, plot_is_static=plot_is_static, width=width, height=height, return_fig=return_fig
     )
     sf.task = "levels"
     return sf
diff --git a/wise_pizza/plotting.py b/wise_pizza/plotting.py
@@ -131,6 +131,7 @@ def plot_segments(
     plot_is_static: bool = False,
     width: int = 2000,
     height: int = 500,
+    return_fig: bool = False
 ):
     """
     Plot segments for explain_levels
@@ -210,7 +211,10 @@ def plot_segments(
             width=width + len(sf.segment_labels) * 30,
         )
     else:
-        fig.show()
+        if return_fig:
+            return fig
+        else:
+            fig.show()
 
 
 def waterfall_args(sf: SliceFinder):