From 8badb6adc1d003dbf0b2b1a800c34821f3fc9aa3 Mon Sep 17 00:00:00 2001
From: wyzula-jan <jan.wyzula@psi.ch>
Date: Fri, 21 Jun 2024 11:56:01 +0200
Subject: [PATCH] fix(cleanup): cleanup added to device_input widgets and
 scan_control

---
 bec_widgets/cli/client.py                     | 19 +++++++++++++++++++
 .../device_combobox/device_combobox.py        |  8 ++++++++
 .../device_inputs/device_input_base.py        |  3 +++
 .../device_line_edit/device_line_edit.py      |  8 ++++++++
 .../widgets/scan_control/scan_control.py      | 17 +++++++++++++++--
 5 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py
index 2428fbe1..72fcf077 100644
--- a/bec_widgets/cli/client.py
+++ b/bec_widgets/cli/client.py
@@ -17,6 +17,7 @@ class Widgets(str, enum.Enum):
     BECDock = "BECDock"
     BECDockArea = "BECDockArea"
     BECFigure = "BECFigure"
+    ScanControl = "ScanControl"
     SpiralProgressBar = "SpiralProgressBar"
     TextBox = "TextBox"
     VSCodeEditor = "VSCodeEditor"
@@ -1823,6 +1824,24 @@ class Ring(RPCBase):
         """
 
 
+class ScanControl(RPCBase):
+    @property
+    @rpc_call
+    def config_dict(self) -> "dict":
+        """
+        Get the configuration of the widget.
+
+        Returns:
+            dict: The configuration of the widget.
+        """
+
+    @rpc_call
+    def get_all_rpc(self) -> "dict":
+        """
+        Get all registered RPC objects.
+        """
+
+
 class SpiralProgressBar(RPCBase):
     @rpc_call
     def get_all_rpc(self) -> "dict":
diff --git a/bec_widgets/widgets/device_inputs/device_combobox/device_combobox.py b/bec_widgets/widgets/device_inputs/device_combobox/device_combobox.py
index 0145908a..f5e4e451 100644
--- a/bec_widgets/widgets/device_inputs/device_combobox/device_combobox.py
+++ b/bec_widgets/widgets/device_inputs/device_combobox/device_combobox.py
@@ -82,3 +82,11 @@ class DeviceComboBox(DeviceInputBase, QComboBox):
         if device_obj is None:
             raise ValueError(f"Device {device_name} is not found.")
         return device_obj
+
+    def cleanup(self):
+        """Cleanup the widget."""
+        super().cleanup()
+
+    def closeEvent(self, event):
+        super().cleanup()
+        QComboBox().closeEvent(event)
\ No newline at end of file
diff --git a/bec_widgets/widgets/device_inputs/device_input_base.py b/bec_widgets/widgets/device_inputs/device_input_base.py
index 7eef2aa2..a250cae3 100644
--- a/bec_widgets/widgets/device_inputs/device_input_base.py
+++ b/bec_widgets/widgets/device_inputs/device_input_base.py
@@ -118,3 +118,6 @@ class DeviceInputBase(BECConnector):
         """
         if device not in self.get_device_list(self.config.device_filter):
             raise ValueError(f"Device {device} is not valid.")
+
+    def cleanup(self):
+        super().cleanup()
diff --git a/bec_widgets/widgets/device_inputs/device_line_edit/device_line_edit.py b/bec_widgets/widgets/device_inputs/device_line_edit/device_line_edit.py
index cf90bfa7..a3756df7 100644
--- a/bec_widgets/widgets/device_inputs/device_line_edit/device_line_edit.py
+++ b/bec_widgets/widgets/device_inputs/device_line_edit/device_line_edit.py
@@ -94,3 +94,11 @@ class DeviceLineEdit(DeviceInputBase, QLineEdit):
         if device_obj is None:
             raise ValueError(f"Device {device_name} is not found.")
         return device_obj
+
+    def cleanup(self):
+        """Cleanup the widget."""
+        super().cleanup()
+
+    def closeEvent(self, event):
+        super().cleanup()
+        QLineEdit().closeEvent(event)
\ No newline at end of file
diff --git a/bec_widgets/widgets/scan_control/scan_control.py b/bec_widgets/widgets/scan_control/scan_control.py
index d450aedc..db21b2d4 100644
--- a/bec_widgets/widgets/scan_control/scan_control.py
+++ b/bec_widgets/widgets/scan_control/scan_control.py
@@ -184,8 +184,21 @@ class ScanControl(BECConnector, QWidget):
         if callable(scan_function):
             scan_function(*args, **kwargs)
 
-    def close(self):
-        super().close()
+    def cleanup(self):
+        self.button_stop_scan.cleanup()
+        if self.arg_box:
+            for widget in self.arg_box.widgets:
+                if hasattr(widget, "cleanup"):
+                    widget.cleanup()
+        for kwarg_box in self.kwarg_boxes:
+            for widget in kwarg_box.widgets:
+                if hasattr(widget, "cleanup"):
+                    widget.cleanup()
+        super().cleanup()
+
+    def closeEvent(self, event):
+        self.cleanup()
+        QWidget().closeEvent(event)
 
 
 # Application example