From ad07bbf85e9c8d9838bdd686f69d41c235b7db19 Mon Sep 17 00:00:00 2001
From: wakonig_k <klaus.wakonig@psi.ch>
Date: Wed, 7 Aug 2024 09:25:32 +0200
Subject: [PATCH] fix(figure): cleanup pyqtgraph

---
 bec_widgets/widgets/figure/figure.py          | 37 +++++++++++--------
 .../widgets/figure/plots/image/image.py       | 12 ++++++
 bec_widgets/widgets/figure/plots/plot_base.py |  8 ++++
 3 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/bec_widgets/widgets/figure/figure.py b/bec_widgets/widgets/figure/figure.py
index 01fe7c96..83138dfc 100644
--- a/bec_widgets/widgets/figure/figure.py
+++ b/bec_widgets/widgets/figure/figure.py
@@ -513,6 +513,13 @@ class BECFigure(BECWidget, pg.GraphicsLayoutWidget):
         if widget_id in self._widgets:
             raise ValueError(f"Widget with ID '{widget_id}' already exists.")
 
+        # Check if position is occupied
+        if row is not None and col is not None:
+            if self.getItem(row, col):
+                raise ValueError(f"Position at row {row} and column {col} is already occupied.")
+        else:
+            row, col = self._find_next_empty_position()
+
         widget = self.widget_handler.create_widget(
             widget_type=widget_type,
             widget_id=widget_id,
@@ -525,23 +532,11 @@ class BECFigure(BECWidget, pg.GraphicsLayoutWidget):
         # used otherwise multiple times
         widget.set_gui_id(widget_id)
 
-        # Check if position is occupied
-        if row is not None and col is not None:
-            if self.getItem(row, col):
-                raise ValueError(f"Position at row {row} and column {col} is already occupied.")
+        widget.config.row = row
+        widget.config.col = col
 
-            widget.config.row = row
-            widget.config.col = col
-
-            # Add widget to the figure
-            self.addItem(widget, row=row, col=col)
-        else:
-            row, col = self._find_next_empty_position()
-            widget.config.row = row
-            widget.config.col = col
-
-            # Add widget to the figure
-            self.addItem(widget, row=row, col=col)
+        # Add widget to the figure
+        self.addItem(widget, row=row, col=col)
 
         # Update num_cols and num_rows based on the added widget
         self.config.num_rows = max(self.config.num_rows, row + 1)
@@ -620,6 +615,7 @@ class BECFigure(BECWidget, pg.GraphicsLayoutWidget):
         """
         if widget_id in self._widgets:
             widget = self._widgets.pop(widget_id)
+            widget.cleanup_pyqtgraph()
             widget.cleanup()
             self.removeItem(widget)
             self.grid[widget.config.row][widget.config.col] = None
@@ -745,3 +741,12 @@ class BECFigure(BECWidget, pg.GraphicsLayoutWidget):
         self.config = FigureConfig(
             widget_class=self.__class__.__name__, gui_id=self.gui_id, theme=theme
         )
+
+    def cleanup_pyqtgraph_all_widgets(self):
+        """Clean up the pyqtgraph widget."""
+        for widget in self.widget_list:
+            widget.cleanup_pyqtgraph()
+
+    def cleanup(self):
+        """Close the figure widget."""
+        self.cleanup_pyqtgraph_all_widgets()
diff --git a/bec_widgets/widgets/figure/plots/image/image.py b/bec_widgets/widgets/figure/plots/image/image.py
index 336e3ee9..1957292f 100644
--- a/bec_widgets/widgets/figure/plots/image/image.py
+++ b/bec_widgets/widgets/figure/plots/image/image.py
@@ -681,3 +681,15 @@ class BECImageShow(BECPlotBase):
                 self.on_image_update, MessageEndpoints.device_monitor_2d(monitor)
             )
         self.images.clear()
+
+    def cleanup_pyqtgraph(self):
+        """Cleanup pyqtgraph items."""
+        super().cleanup_pyqtgraph()
+        item = self.plot_item
+        cbar = item.items[0].color_bar
+        cbar.vb.menu.close()
+        cbar.vb.menu.deleteLater()
+        cbar.gradient.menu.close()
+        cbar.gradient.menu.deleteLater()
+        cbar.gradient.colorDialog.close()
+        cbar.gradient.colorDialog.deleteLater()
diff --git a/bec_widgets/widgets/figure/plots/plot_base.py b/bec_widgets/widgets/figure/plots/plot_base.py
index 427852bc..0b130be7 100644
--- a/bec_widgets/widgets/figure/plots/plot_base.py
+++ b/bec_widgets/widgets/figure/plots/plot_base.py
@@ -314,3 +314,11 @@ class BECPlotBase(BECConnector, pg.GraphicsLayout):
         """Remove the plot widget from the figure."""
         if self.figure is not None:
             self.figure.remove(widget_id=self.gui_id)
+
+    def cleanup_pyqtgraph(self):
+        """Cleanup pyqtgraph items."""
+        item = self.plot_item
+        item.vb.menu.close()
+        item.vb.menu.deleteLater()
+        item.ctrlMenu.close()
+        item.ctrlMenu.deleteLater()