Merge branch '256_plot_finalize' into 'dev'

Shared finalize API for all return_fig plot methods

See merge request mass-spectrometry/corems!239

Author: kheal

corems/chroma_peak/factory/chroma_peak_classes.py

@@ -10,6 +10,7 @@
     GCPeakCalculation,
     LCMSMassFeatureCalculation,
 )
+from corems.encapsulation.plot_utils import _finalize_plot
 from corems.mass_spectra.factory.chromat_data import EIC_Data
 from corems.molecular_id.factory.EI_SQL import LowResCompoundRef
 
@@ -594,6 +595,7 @@ def plot(
         self,
         to_plot=["EIC", "MS1", "MS2"],
         return_fig=True,
+        path=None,
         plot_smoothed_eic=False,
         plot_eic_datapoints=False,
         molecular_metadata=None,
@@ -608,7 +610,11 @@ def plot(
             "EIC", "MS2", "MS2_mirror", and "MS1".
             Default is ["EIC", "MS1", "MS2"].
         return_fig : bool, optional
-            If True, the figure is returned. Default is True.
+            If True, return the open figure (caller owns lifecycle).
+            Default is True.
+        path : str or path-like, optional
+            If set, save the figure to this path. When ``return_fig`` is False,
+            the figure is closed after saving and ``plt.show()`` is not called.
         plot_smoothed_eic : bool, optional
             If True, the smoothed EIC is plotted. Default is False.
         plot_eic_datapoints : bool, optional
@@ -623,8 +629,9 @@ def plot(
         Returns
         -------
         matplotlib.figure.Figure or None
-            The figure object if `return_fig` is True.
-            Otherwise None and the figure is displayed.
+            The figure object if `return_fig` is True (left open for the caller).
+            Otherwise None; the figure is displayed with ``plt.show()`` when
+            ``path`` is not set.
         """
         # Adjust to_plot list if there are not spectra added to the mass features
         if self.mass_spectrum is None:
@@ -681,10 +688,7 @@ def plot(
         # Add space between subplots
         plt.tight_layout()
 
-        if return_fig:
-            # Close figure
-            plt.close(fig)
-            return fig
+        return _finalize_plot(fig, return_fig=return_fig, path=path)
 
     @property
     def mz(self):

corems/encapsulation/plot_utils.py

@@ -0,0 +1,51 @@
+"""Shared plotting helpers for CoreMS figure lifecycle management."""
+
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+
+
+def _finalize_plot(fig, return_fig=False, path=None, **savefig_kwargs):
+    """Shared exit path for CoreMS plot methods that support ``return_fig``.
+
+    Parameters
+    ----------
+    fig : matplotlib.figure.Figure
+        Figure to finalize.
+    return_fig : bool, optional
+        If True, leave the figure open and return it. The caller owns the
+        figure lifecycle (e.g. further customization or ``plt.close(fig)``).
+        Default is False.
+    path : str or path-like, optional
+        If set, save the figure to this path via ``fig.savefig`` before
+        showing or returning.
+    **savefig_kwargs
+        Forwarded to ``fig.savefig`` when ``path`` is set.
+
+    Returns
+    -------
+    matplotlib.figure.Figure or None
+        The open figure if ``return_fig`` is True; otherwise None.
+
+    Notes
+    -----
+    Behavior matrix:
+
+    - ``return_fig=False``, ``path=None``: call ``plt.show()`` (interactive /
+      notebook default).
+    - ``return_fig=False``, ``path`` set: save, close the figure, do **not**
+      call ``plt.show()`` (batch / headless friendly).
+    - ``return_fig=True``: optionally save if ``path`` is set; return the open
+      figure without showing or closing it.
+    """
+    if path is not None:
+        fig.savefig(path, **savefig_kwargs)
+
+    if return_fig:
+        return fig
+
+    if path is None:
+        plt.show()
+    else:
+        plt.close(fig)
+    return None