From f73e566f4ff4f056036e4bd7d0c2a48056c4e86b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= <cr52@protonmail.com>
Date: Thu, 5 Dec 2024 17:15:57 +0100
Subject: [PATCH] BUG: make Plotable.plot return unconditionally

---
 nonos/api/analysis.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/nonos/api/analysis.py b/nonos/api/analysis.py
index 20167dca..34aa0fd5 100644
--- a/nonos/api/analysis.py
+++ b/nonos/api/analysis.py
@@ -24,6 +24,7 @@
 from nonos.logging import logger
 
 if TYPE_CHECKING:
+    from matplotlib.artist import Artist
     from matplotlib.axes import Axes
     from matplotlib.figure import Figure
 
@@ -53,7 +54,7 @@ def plot(
         unit_conversion=None,
         nbin=None,  # deprecated
         **kwargs,
-    ):
+    ) -> "Artist":
         if nbin is not None:
             warnings.warn(
                 "The nbin parameter has no effect and is deprecated",
@@ -65,6 +66,7 @@ def plot(
         if log:
             data = np.log10(data)
 
+        artist: Artist
         if self.dimension == 2:
             self.akey = self.dict_plotable["abscissa"]
             self.okey = self.dict_plotable["ordinate"]
@@ -81,7 +83,7 @@ def plot(
                 vmax = kwargs.pop("vmax") if "vmax" in kwargs else np.nanmax(data)
                 kw.update({"vmin": vmin, "vmax": vmax})
 
-            im = ax.pcolormesh(
+            artist = im = ax.pcolormesh(
                 self.avalue,
                 self.ovalue,
                 data,
@@ -120,8 +122,6 @@ def plot(
                             trf, subs=list(range(1, int(trf.base)))
                         )
                         cb_axis.set_minor_locator(locator)
-            else:
-                return im
         elif self.dimension == 1:
             vmin = kwargs.pop("vmin") if "vmin" in kwargs else np.nanmin(data)
             vmax = kwargs.pop("vmax") if "vmax" in kwargs else np.nanmax(data)
@@ -130,7 +130,7 @@ def plot(
             if "norm" in kwargs:
                 logger.info("norm has no meaning in 1D.")
                 kwargs.pop("norm")
-            ax.plot(self.avalue, data, **kwargs)
+            artist = ax.plot(self.avalue, data, **kwargs)[0]
             ax.set_ylim(ymin=vmin)
             ax.set_ylim(ymax=vmax)
             ax.set_xlabel(self.akey)
@@ -144,6 +144,8 @@ def plot(
         if filename is not None:
             fig.savefig(f"{filename}.{fmt}", bbox_inches="tight", dpi=dpi)
 
+        return artist
+
 
 class Coordinates:
     """Coordinates class from x1, x2, x3"""