Merge pull request #208 from fooof-tools/savef

[ENH] - Extend access to options in savefig
fooof-tools · May 5, 2021 · c19a41a · c19a41a
2 parents 62196db + 9cc9a30
commit c19a41a
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git a/fooof/plts/utils.py b/fooof/plts/utils.py
@@ -189,11 +189,21 @@ def decorated(*args, **kwargs):
         #   Defaults to saving when file name given (since bool(str)->True; bool(None)->False)
         save_fig = kwargs.pop('save_fig', bool(file_name))
 
+        # Check any collect any other plot keywords
+        save_kwargs = kwargs.pop('save_kwargs', {})
+        save_kwargs.setdefault('bbox_inches', 'tight')
+
+        # Check and collect whether to close the plot
+        close = kwargs.pop('close', None)
+
         func(*args, **kwargs)
 
         if save_fig:
             if not file_name:
                 raise ValueError("Input 'file_name' is required to save out the plot.")
-            plt.savefig(fpath(file_path, fname(file_name, 'png')))
+            plt.savefig(fpath(file_path, fname(file_name, 'png')), **save_kwargs)
+
+        if close:
+            plt.close()
 
     return decorated
diff --git a/fooof/tests/plts/test_utils.py b/fooof/tests/plts/test_utils.py
@@ -87,6 +87,11 @@ def example_plot():
     example_plot(save_fig=True, file_path=TEST_PLOTS_PATH, file_name='test_savefig2.pdf')
     assert os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_savefig2.pdf'))
 
+    # Test giving additional save kwargs
+    example_plot(file_path=TEST_PLOTS_PATH, file_name='test_savefig3.pdf',
+                 save_kwargs={'facecolor' : 'red'})
+    assert os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_savefig3.pdf'))
+
     # Test does not save when `save_fig` set to False
-    example_plot(save_fig=False, file_path=TEST_PLOTS_PATH, file_name='test_savefig3.pdf')
-    assert not os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_savefig3.pdf'))
+    example_plot(save_fig=False, file_path=TEST_PLOTS_PATH, file_name='test_savefig_nope.pdf')
+    assert not os.path.exists(os.path.join(TEST_PLOTS_PATH, 'test_savefig_nope.pdf'))