bec-widgets 0.69.0__py3-none-any.whl → 0.71.0__py3-none-any.whl

bec_widgets/widgets/scan_control/scan_control.py

@@ -1,53 +1,37 @@
+import qdarktheme
 from bec_lib.endpoints import MessageEndpoints
 from qtpy.QtWidgets import (
     QApplication,
-    QCheckBox,
     QComboBox,
-    QDoubleSpinBox,
-    QFrame,
     QGridLayout,
     QGroupBox,
-    QHBoxLayout,
-    QHeaderView,
-    QLabel,
-    QLayout,
-    QLineEdit,
     QPushButton,
-    QSpinBox,
-    QTableWidget,
-    QTableWidgetItem,
+    QSizePolicy,
     QVBoxLayout,
     QWidget,
 )
 
-from bec_widgets.utils.bec_dispatcher import BECDispatcher
-from bec_widgets.utils.widget_io import WidgetIO
+from bec_widgets.utils import BECConnector
+from bec_widgets.widgets.buttons.stop_button.stop_button import StopButton
+from bec_widgets.widgets.scan_control.scan_group_box import ScanGroupBox
 
 
-class ScanArgType:
-    DEVICE = "device"
-    FLOAT = "float"
-    INT = "int"
-    BOOL = "bool"
-    STR = "str"
+class ScanControl(BECConnector, QWidget):
 
-
-class ScanControl(QWidget):
-    WIDGET_HANDLER = {
-        ScanArgType.DEVICE: QLineEdit,
-        ScanArgType.FLOAT: QDoubleSpinBox,
-        ScanArgType.INT: QSpinBox,
-        ScanArgType.BOOL: QCheckBox,
-        ScanArgType.STR: QLineEdit,
-    }
-
-    def __init__(self, parent=None, client=None, allowed_scans=None):
-        super().__init__(parent)
+    def __init__(
+        self, parent=None, client=None, gui_id: str | None = None, allowed_scans: list | None = None
+    ):
+        super().__init__(client=client, gui_id=gui_id)
+        QWidget.__init__(self, parent=parent)
 
         # Client from BEC + shortcuts to device manager and scans
-        self.client = BECDispatcher().client if client is None else client
-        self.dev = self.client.device_manager.devices
-        self.scans = self.client.scans
+        self.get_bec_shortcuts()
+
+        # Main layout
+        self.layout = QVBoxLayout(self)
+        self.arg_box = None
+        self.kwarg_boxes = []
+        self.expert_mode = False  # TODO implement in the future versions
 
         # Scan list - allowed scans for the GUI
         self.allowed_scans = allowed_scans
@@ -56,389 +40,173 @@ class ScanControl(QWidget):
         self._init_UI()
 
     def _init_UI(self):
-        self.verticalLayout = QVBoxLayout(self)
+        """
+        Initializes the UI of the scan control widget. Create the top box for scan selection and populate scans to main combobox.
+        """
 
         # Scan selection group box
-        self.scan_selection_group = QGroupBox("Scan Selection", self)
-        self.scan_selection_layout = QVBoxLayout(self.scan_selection_group)
-        self.comboBox_scan_selection = QComboBox(self.scan_selection_group)
-        self.button_run_scan = QPushButton("Run Scan", self.scan_selection_group)
-        self.scan_selection_layout.addWidget(self.comboBox_scan_selection)
-        self.scan_selection_layout.addWidget(self.button_run_scan)
-        self.verticalLayout.addWidget(self.scan_selection_group)
-
-        # Scan control group box
-        self.scan_control_group = QGroupBox("Scan Control", self)
-        self.scan_control_layout = QVBoxLayout(self.scan_control_group)
-        self.verticalLayout.addWidget(self.scan_control_group)
-
-        # Kwargs layout - just placeholder
-        self.kwargs_layout = QGridLayout()
-        self.scan_control_layout.addLayout(self.kwargs_layout)
-
-        # 1st Separator
-        self.add_horizontal_separator(self.scan_control_layout)
-
-        # Buttons
-        self.button_layout = QHBoxLayout()
-        self.pushButton_add_bundle = QPushButton("Add Bundle", self.scan_control_group)
-        self.pushButton_add_bundle.clicked.connect(self.add_bundle)
-        self.pushButton_remove_bundle = QPushButton("Remove Bundle", self.scan_control_group)
-        self.pushButton_remove_bundle.clicked.connect(self.remove_bundle)
-        self.button_layout.addWidget(self.pushButton_add_bundle)
-        self.button_layout.addWidget(self.pushButton_remove_bundle)
-        self.scan_control_layout.addLayout(self.button_layout)
-
-        # 2nd Separator
-        self.add_horizontal_separator(self.scan_control_layout)
-
-        # Initialize the QTableWidget for args
-        self.args_table = QTableWidget()
-        self.args_table.verticalHeader().setSectionResizeMode(QHeaderView.Fixed)
-
-        self.scan_control_layout.addWidget(self.args_table)
+        self.scan_selection_group = self.create_scan_selection_group()
+        self.scan_selection_group.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Fixed)
+        self.layout.addWidget(self.scan_selection_group)
 
         # Connect signals
         self.comboBox_scan_selection.currentIndexChanged.connect(self.on_scan_selected)
         self.button_run_scan.clicked.connect(self.run_scan)
+        self.button_add_bundle.clicked.connect(self.add_arg_bundle)
+        self.button_remove_bundle.clicked.connect(self.remove_arg_bundle)
 
         # Initialize scan selection
         self.populate_scans()
 
-    def add_horizontal_separator(self, layout) -> None:
+    def create_scan_selection_group(self) -> QGroupBox:
         """
-        Adds a horizontal separator to the given layout
+        Creates the scan selection group box with combobox to select the scan and start/stop button.
 
-        Args:
-            layout: Layout to add the separator to
-        """
-        separator = QFrame(self.scan_control_group)
-        separator.setFrameShape(QFrame.HLine)
-        separator.setFrameShadow(QFrame.Sunken)
-        layout.addWidget(separator)
+        Returns:
+            QGroupBox: Group box containing the scan selection widgets.
+        """
+
+        scan_selection_group = QGroupBox("Scan Selection", self)
+        self.scan_selection_layout = QGridLayout(scan_selection_group)
+        self.comboBox_scan_selection = QComboBox(scan_selection_group)
+        # Run button
+        self.button_run_scan = QPushButton("Start", scan_selection_group)
+        self.button_run_scan.setStyleSheet("background-color:  #559900; color: white")
+        # Stop button
+        self.button_stop_scan = StopButton(parent=scan_selection_group)
+        # Add bundle button
+        self.button_add_bundle = QPushButton("Add Bundle", scan_selection_group)
+        # Remove bundle button
+        self.button_remove_bundle = QPushButton("Remove Bundle", scan_selection_group)
+
+        self.scan_selection_layout.addWidget(self.comboBox_scan_selection, 0, 0, 1, 2)
+        self.scan_selection_layout.addWidget(self.button_run_scan, 1, 0)
+        self.scan_selection_layout.addWidget(self.button_stop_scan, 1, 1)
+        self.scan_selection_layout.addWidget(self.button_add_bundle, 2, 0)
+        self.scan_selection_layout.addWidget(self.button_remove_bundle, 2, 1)
+
+        return scan_selection_group
 
     def populate_scans(self):
-        """Populates the scan selection combo box with available scans"""
-        self.available_scans = self.client.producer.get(MessageEndpoints.available_scans()).resource
+        """Populates the scan selection combo box with available scans from BEC session."""
+        self.available_scans = self.client.connector.get(
+            MessageEndpoints.available_scans()
+        ).resource
         if self.allowed_scans is None:
-            allowed_scans = self.available_scans.keys()
+            supported_scans = ["ScanBase", "SyncFlyScanBase", "AsyncFlyScanBase"]
+            allowed_scans = [
+                scan_name
+                for scan_name, scan_info in self.available_scans.items()
+                if scan_info["base_class"] in supported_scans and len(scan_info["gui_config"]) > 0
+            ]
+
         else:
             allowed_scans = self.allowed_scans
-        # TODO check parent class is ScanBase -> filter out the scans not relevant for GUI
         self.comboBox_scan_selection.addItems(allowed_scans)
 
     def on_scan_selected(self):
         """Callback for scan selection combo box"""
+        self.reset_layout()
         selected_scan_name = self.comboBox_scan_selection.currentText()
         selected_scan_info = self.available_scans.get(selected_scan_name, {})
 
-        print(selected_scan_info)  # TODO remove when widget will be more mature
-        # Generate kwargs input
-        self.generate_kwargs_input_fields(selected_scan_info)
-
-        # Args section
-        self.generate_args_input_fields(selected_scan_info)
-
-    def add_labels_to_layout(self, labels: list, grid_layout: QGridLayout) -> None:
-        """
-        Adds labels to the given grid layout as a separate row.
-
-        Args:
-            labels (list): List of label names to add.
-            grid_layout (QGridLayout): The grid layout to which labels will be added.
-        """
-        row_index = grid_layout.rowCount()  # Get the next available row
-        for column_index, label_name in enumerate(labels):
-            label = QLabel(label_name.capitalize(), self.scan_control_group)
-            # Add the label to the grid layout at the calculated row and current column
-            grid_layout.addWidget(label, row_index, column_index)
-
-    def add_labels_to_table(
-        self, labels: list, table: QTableWidget
-    ) -> None:  # TODO could be moved to BECTable
-        """
-        Adds labels to the given table widget as a header row.
-
-        Args:
-            labels(list): List of label names to add.
-            table(QTableWidget): The table widget to which labels will be added.
-        """
-        table.setColumnCount(len(labels))
-        table.setHorizontalHeaderLabels(labels)
-
-    def generate_args_input_fields(self, scan_info: dict) -> None:
-        """
-        Generates input fields for args.
+        gui_config = selected_scan_info.get("gui_config", {})
+        self.arg_group = gui_config.get("arg_group", None)
+        self.kwarg_groups = gui_config.get("kwarg_groups", None)
 
-        Args:
-            scan_info(dict): Scan signature dictionary from BEC.
-        """
+        if self.arg_box is None:
+            self.button_add_bundle.setEnabled(False)
+            self.button_remove_bundle.setEnabled(False)
 
-        # Setup args table limits
-        self.set_args_table_limits(self.args_table, scan_info)
+        if len(self.arg_group["arg_inputs"]) > 0:
+            self.button_add_bundle.setEnabled(True)
+            self.button_remove_bundle.setEnabled(True)
+            self.add_arg_group(self.arg_group)
+        if len(self.kwarg_groups) > 0:
+            self.add_kwargs_boxes(self.kwarg_groups)
 
-        # Get arg_input from selected scan
-        self.arg_input = scan_info.get("arg_input", {})
+        self.update()
+        self.adjustSize()
 
-        # Generate labels for table
-        self.add_labels_to_table(list(self.arg_input.keys()), self.args_table)
-
-        # add minimum number of args rows
-        if self.arg_size_min is not None:
-            for i in range(self.arg_size_min):
-                self.add_bundle()
-
-    def generate_kwargs_input_fields(self, scan_info: dict) -> None:
+    def add_kwargs_boxes(self, groups: list):
         """
-        Generates input fields for kwargs
+        Adds the given gui_groups to the scan control layout.
 
         Args:
-            scan_info(dict): Scan signature dictionary from BEC.
-        """
-        # Create a new kwarg layout to replace the old one - this is necessary because otherwise row count is not reseted
-        self.clear_and_delete_layout(self.kwargs_layout)
-        self.kwargs_layout = self.create_new_grid_layout()  # Create new grid layout
-        self.scan_control_layout.insertLayout(0, self.kwargs_layout)
-
-        # Get signature
-        signature = scan_info.get("signature", [])
-
-        # Extract kwargs from the converted signature
-        kwargs = [param["name"] for param in signature if param["kind"] == "KEYWORD_ONLY"]
-
-        # Add labels
-        self.add_labels_to_layout(kwargs, self.kwargs_layout)
-
-        # Add widgets
-        widgets = self.generate_widgets_from_signature(kwargs, signature)
-
-        self.add_widgets_row_to_layout(self.kwargs_layout, widgets)
-
-    def generate_widgets_from_signature(self, items: list, signature: dict = None) -> list:
+            groups(list): List of dictionaries containing the gui_group information.
         """
-        Generates widgets from the given list of items.
-
-        Args:
-            items(list): List of items to create widgets for.
-            signature(dict, optional): Scan signature dictionary from BEC.
+        for group in groups:
+            box = ScanGroupBox(box_type="kwargs", config=group)
+            box.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Fixed)
+            self.layout.addWidget(box)
+            self.kwarg_boxes.append(box)
 
-        Returns:
-            list: List of widgets created from the given items.
-        """
-        widgets = []  # Initialize an empty list to hold the widgets
-
-        for item in items:
-            if signature:
-                # If a signature is provided, extract type and name from it
-                kwarg_info = next((info for info in signature if info["name"] == item), None)
-                if kwarg_info:
-                    item_type = kwarg_info.get("annotation", "_empty")
-                    item_name = item
-            else:
-                # If no signature is provided, assume the item is a tuple of (name, type)
-                item_name, item_type = item
-
-            widget_class = self.WIDGET_HANDLER.get(item_type, None)
-            if widget_class is None:
-                print(f"Unsupported annotation '{item_type}' for parameter '{item_name}'")
-                continue
-
-            # Instantiate the widget and set some properties if necessary
-            widget = widget_class()
-
-            # set high default range for spin boxes #TODO can be linked to motor/device limits from BEC
-            if isinstance(widget, (QSpinBox, QDoubleSpinBox)):
-                widget.setRange(-9999, 9999)
-                widget.setValue(0)
-            # Add the widget to the list
-            widgets.append(widget)
-
-        return widgets
-
-    def set_args_table_limits(self, table: QTableWidget, scan_info: dict) -> None:
-        # Get bundle info
-        arg_bundle_size = scan_info.get("arg_bundle_size", {})
-        self.arg_size_min = arg_bundle_size.get("min", 1)
-        self.arg_size_max = arg_bundle_size.get("max", None)
-
-        # Clear the previous input fields
-        table.setRowCount(0)  # Wipe table
-
-    def add_widgets_row_to_layout(
-        self, grid_layout: QGridLayout, widgets: list, row_index: int = None
-    ) -> None:
+    def add_arg_group(self, group: dict):
         """
-        Adds a row of widgets to the given grid layout.
+        Adds the given gui_groups to the scan control layout.
 
         Args:
-            grid_layout (QGridLayout): The grid layout to which widgets will be added.
-            items (list): List of parameter names to create widgets for.
-            row_index (int): The row index where the widgets should be added.
         """
-        # If row_index is not specified, add to the next available row
-        if row_index is None:
-            row_index = grid_layout.rowCount()
+        self.arg_box = ScanGroupBox(box_type="args", config=group)
+        self.arg_box.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Fixed)
+        self.layout.addWidget(self.arg_box)
 
-        for column_index, widget in enumerate(widgets):
-            # Add the widget to the grid layout at the specified row and column
-            grid_layout.addWidget(widget, row_index, column_index)
-
-    def add_widgets_row_to_table(
-        self, table_widget: QTableWidget, widgets: list, row_index: int = None
-    ) -> None:
-        """
-        Adds a row of widgets to the given QTableWidget.
+    def add_arg_bundle(self):
+        self.arg_box.add_widget_bundle()
 
-        Args:
-            table_widget (QTableWidget): The table widget to which widgets will be added.
-            widgets (list): List of widgets to add to the table.
-            row_index (int): The row index where the widgets should be added. If None, add to the end.
-        """
-        # If row_index is not specified, add to the end of the table
-        if row_index is None or row_index > table_widget.rowCount():
-            row_index = table_widget.rowCount()
-            if self.arg_size_max is not None:  # ensure the max args size is not exceeded
-                if row_index >= self.arg_size_max:
-                    return
-            table_widget.insertRow(row_index)
-
-        for column_index, widget in enumerate(widgets):
-            # If the widget is a subclass of QWidget, use setCellWidget
-            if issubclass(type(widget), QWidget):
-                table_widget.setCellWidget(row_index, column_index, widget)
-            else:
-                # Otherwise, assume it's a string or some other value that should be displayed as text
-                item = QTableWidgetItem(str(widget))
-                table_widget.setItem(row_index, column_index, item)
-
-        # Optionally, adjust the row height based on the content #TODO decide if needed
-        table_widget.setRowHeight(
-            row_index,
-            max(widget.sizeHint().height() for widget in widgets if isinstance(widget, QWidget)),
-        )
-
-    def remove_last_row_from_table(self, table_widget: QTableWidget) -> None:
-        """
-        Removes the last row from the given QTableWidget until only one row is left.
-
-        Args:
-            table_widget (QTableWidget): The table widget from which the last row will be removed.
-        """
-        row_count = table_widget.rowCount()
-        if (
-            row_count > self.arg_size_min
-        ):  # Check to ensure there is a minimum number of rows remaining
-            table_widget.removeRow(row_count - 1)
-
-    def create_new_grid_layout(self):
-        new_layout = QGridLayout()
-        # TODO maybe setup other layouts properties here?
-        return new_layout
-
-    def clear_and_delete_layout(self, layout: QLayout):
-        """
-        Clears and deletes the given layout and all its child widgets.
+    def remove_arg_bundle(self):
+        self.arg_box.remove_widget_bundle()
 
-        Args:
-            layout(QLayout): Layout to clear and delete
-        """
-        if layout is not None:
-            while layout.count():
-                item = layout.takeAt(0)
-                widget = item.widget()
-                if widget:
-                    widget.deleteLater()
-                else:
-                    sub_layout = item.layout()
-                    if sub_layout:
-                        self.clear_and_delete_layout(sub_layout)
-            layout.deleteLater()
-
-    def add_bundle(self) -> None:
-        """Adds a new bundle to the scan control layout"""
-        # Get widgets used for particular scan and save them to be able to use for adding bundles
-        args_widgets = self.generate_widgets_from_signature(
-            self.arg_input.items()
-        )  # TODO decide if make sense to put widget list into method parameters
-
-        # Add first widgets row to the table
-        self.add_widgets_row_to_table(self.args_table, args_widgets)
-
-    def remove_bundle(self) -> None:
-        """Removes the last bundle from the scan control layout"""
-        self.remove_last_row_from_table(self.args_table)
-
-    def extract_kwargs_from_grid_row(self, grid_layout: QGridLayout, row: int) -> dict:
-        kwargs = {}
-        for column in range(grid_layout.columnCount()):
-            label_item = grid_layout.itemAtPosition(row, column)
-            if label_item is not None:
-                label_widget = label_item.widget()
-                if isinstance(label_widget, QLabel):
-                    key = label_widget.text()
-
-                    # The corresponding value widget is in the next row
-                    value_item = grid_layout.itemAtPosition(row + 1, column)
-                    if value_item is not None:
-                        value_widget = value_item.widget()
-                        # Use WidgetIO.get_value to extract the value
-                        value = WidgetIO.get_value(value_widget)
-                        kwargs[key] = value
-        return kwargs
-
-    def extract_args_from_table(self, table: QTableWidget) -> list:
-        """
-        Extracts the arguments from the given table widget.
+    def reset_layout(self):
+        """Clears the scan control layout from GuiGroups and ArgGroups boxes."""
+        if self.arg_box is not None:
+            self.layout.removeWidget(self.arg_box)
+            self.arg_box.deleteLater()
+            self.arg_box = None
+        if self.kwarg_boxes != []:
+            self.remove_kwarg_boxes()
 
-        Args:
-            table(QTableWidget): Table widget from which to extract the arguments
-        """
-        args = []
-        for row in range(table.rowCount()):
-            row_args = []
-            for column in range(table.columnCount()):
-                widget = table.cellWidget(row, column)
-                if widget:
-                    if isinstance(widget, QLineEdit):  # special case for QLineEdit for Devices
-                        value = widget.text().lower()
-                        if value in self.dev:
-                            value = getattr(self.dev, value)
-                        else:
-                            raise ValueError(f"The device '{value}' is not recognized.")
-                    else:
-                        value = WidgetIO.get_value(widget)
-                    row_args.append(value)
-            args.extend(row_args)
-        return args
+    def remove_kwarg_boxes(self):
+        for box in self.kwarg_boxes:
+            self.layout.removeWidget(box)
+            box.deleteLater()
+        self.kwarg_boxes = []
 
     def run_scan(self):
-        # Extract kwargs for the scan
-        kwargs = {
-            k.lower(): v
-            for k, v in self.extract_kwargs_from_grid_row(self.kwargs_layout, 1).items()
-        }
-
-        # Extract args from the table
-        args = self.extract_args_from_table(self.args_table)
-
-        # Convert args to lowercase if they are strings
-        args = [arg.lower() if isinstance(arg, str) else arg for arg in args]
-
-        # Execute the scan
+        args = []
+        kwargs = {}
+        if self.arg_box is not None:
+            args = self.arg_box.get_parameters()
+        for box in self.kwarg_boxes:
+            box_kwargs = box.get_parameters()
+            kwargs.update(box_kwargs)
         scan_function = getattr(self.scans, self.comboBox_scan_selection.currentText())
         if callable(scan_function):
             scan_function(*args, **kwargs)
 
+    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
 if __name__ == "__main__":  # pragma: no cover
-    # BECclient global variables
-    client = BECDispatcher().client
-    client.start()
-
     app = QApplication([])
-    scan_control = ScanControl(client=client)  # allowed_scans=["line_scan", "grid_scan"])
+    scan_control = ScanControl()
 
+    qdarktheme.setup_theme("auto")
     window = scan_control
     window.show()
     app.exec()