From f2e5a85e616aa76d4b7ad3b3c76a24ba114ebdd1 Mon Sep 17 00:00:00 2001 From: wakonig_k Date: Thu, 14 Aug 2025 22:08:51 +0200 Subject: [PATCH] feat(scan control): add support for literals --- .../control/scan_control/scan_group_box.py | 49 +++++++++++++++++-- tests/unit_tests/test_scan_control.py | 16 +++++- 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/bec_widgets/widgets/control/scan_control/scan_group_box.py b/bec_widgets/widgets/control/scan_control/scan_group_box.py index 7d0d3826..63a24609 100644 --- a/bec_widgets/widgets/control/scan_control/scan_group_box.py +++ b/bec_widgets/widgets/control/scan_control/scan_group_box.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, Sequence from bec_lib.logger import bec_logger from bec_qthemes import material_icon @@ -36,7 +36,7 @@ class ScanArgType: BOOL = "bool" STR = "str" DEVICEBASE = "DeviceBase" - LITERALS = "dict" + LITERALS_DICT = "dict" # Used when the type is provided as a dict with Literal key class SettingsDialog(QDialog): @@ -83,6 +83,39 @@ class ScanSpinBox(QSpinBox): self.setValue(default) +class ScanLiteralsComboBox(QComboBox): + def __init__( + self, parent=None, arg_name: str | None = None, default: str | None = None, *args, **kwargs + ): + super().__init__(parent=parent, *args, **kwargs) + self.arg_name = arg_name + self.default = default + if default is not None: + self.setCurrentText(default) + + def set_literals(self, literals: Sequence[str | int | float | None]) -> None: + """ + Set the list of literals for the combo box. + + Args: + literals: List of literal values (can be strings, integers, floats or None) + """ + self.clear() + literals = set(literals) # Remove duplicates + if None in literals: + literals.remove(None) + self.addItem("") + + self.addItems([str(value) for value in literals]) + + # find index of the default value + index = max(self.findText(str(self.default)), 0) + self.setCurrentIndex(index) + + def get_value(self) -> str | None: + return self.currentText() if self.currentText() else None + + class ScanDoubleSpinBox(QDoubleSpinBox): def __init__( self, parent=None, arg_name: str = None, default: float | None = None, *args, **kwargs @@ -137,7 +170,7 @@ class ScanGroupBox(QGroupBox): ScanArgType.INT: ScanSpinBox, ScanArgType.BOOL: ScanCheckBox, ScanArgType.STR: ScanLineEdit, - ScanArgType.LITERALS: QComboBox, # TODO figure out combobox logic + ScanArgType.LITERALS_DICT: ScanLiteralsComboBox, } device_selected = Signal(str) @@ -226,7 +259,11 @@ class ScanGroupBox(QGroupBox): for column_index, item in enumerate(group_inputs): arg_name = item.get("name", None) default = item.get("default", None) - widget_class = self.WIDGET_HANDLER.get(item["type"], None) + item_type = item.get("type", None) + if isinstance(item_type, dict) and "Literal" in item_type: + widget_class = self.WIDGET_HANDLER.get(ScanArgType.LITERALS_DICT, None) + else: + widget_class = self.WIDGET_HANDLER.get(item["type"], None) if widget_class is None: logger.error( f"Unsupported annotation '{item['type']}' for parameter '{item['name']}'" @@ -239,6 +276,8 @@ class ScanGroupBox(QGroupBox): widget.set_device_filter(BECDeviceFilter.DEVICE) self.selected_devices[widget] = "" widget.device_selected.connect(self.emit_device_selected) + if isinstance(widget, ScanLiteralsComboBox): + widget.set_literals(item["type"].get("Literal", [])) tooltip = item.get("tooltip", None) if tooltip is not None: widget.setToolTip(item["tooltip"]) @@ -336,6 +375,8 @@ class ScanGroupBox(QGroupBox): widget = self.layout.itemAtPosition(1, i).widget() if isinstance(widget, DeviceLineEdit) and device_object: value = widget.get_current_device().name + elif isinstance(widget, ScanLiteralsComboBox): + value = widget.get_value() else: value = WidgetIO.get_value(widget) kwargs[widget.arg_name] = value diff --git a/tests/unit_tests/test_scan_control.py b/tests/unit_tests/test_scan_control.py index 0d0c24b3..30ce317d 100644 --- a/tests/unit_tests/test_scan_control.py +++ b/tests/unit_tests/test_scan_control.py @@ -210,6 +210,15 @@ available_scans_message = AvailableResourceMessage( "default": False, "expert": False, }, + { + "arg": False, + "name": "optim_trajectory", + "type": {"Literal": ("option1", "option2", "option3", None)}, + "display_name": "Optim Trajectory", + "tooltip": None, + "default": None, + "expert": False, + }, ], } ], @@ -304,7 +313,10 @@ def test_on_scan_selected(scan_control, scan_name): label = kwarg_box.layout.itemAtPosition(0, index).widget() assert label.text() == kwarg_info["display_name"] widget = kwarg_box.layout.itemAtPosition(1, index).widget() - expected_widget_type = kwarg_box.WIDGET_HANDLER.get(kwarg_info["type"], None) + if isinstance(kwarg_info["type"], dict) and "Literal" in kwarg_info["type"]: + expected_widget_type = kwarg_box.WIDGET_HANDLER.get("dict", None) + else: + expected_widget_type = kwarg_box.WIDGET_HANDLER.get(kwarg_info["type"], None) assert isinstance(widget, expected_widget_type) @@ -441,7 +453,7 @@ def test_run_grid_scan_with_parameters(scan_control, mocked_client): args_row2["steps"], ] assert called_args == tuple(expected_args_list) - assert called_kwargs == kwargs | {"metadata": {"sample_name": ""}} + assert called_kwargs == kwargs | {"metadata": {"sample_name": ""}, "optim_trajectory": None} # Check the emitted signal mock_slot.assert_called_once()