wip sorting list

bec_widgets/utils/list_of_expandable_frames.py

@@ -5,8 +5,9 @@ from typing import Generic, Iterable, NamedTuple, TypeVar
 
 from bec_lib.logger import bec_logger
 from more_itertools import consume
+from PySide6.QtGui import QFont
 from PySide6.QtWidgets import QListWidgetItem, QWidget
-from qtpy.QtCore import QSize
+from qtpy.QtCore import QSize, Qt
 from qtpy.QtWidgets import QListWidget
 
 from bec_widgets.utils.error_popups import SafeSlot
@@ -14,6 +15,8 @@ from bec_widgets.utils.expandable_frame import ExpandableGroupFrame
 
 logger = bec_logger.logger
 
+_SORT_KEY_ROLE = 117
+
 _EF = TypeVar("_EF", bound=ExpandableGroupFrame)
 
 
@@ -54,19 +57,31 @@ class ListOfExpandableFrames(QListWidget, Generic[_EF]):
             item.setSizeHint(QSize(item_widget.width(), item_widget.height()))
 
         item = QListWidgetItem(self)
-        item_widget = self._item_class(*args, **kwargs)
+        item.setData(_SORT_KEY_ROLE, id)  # used for sorting
 
+        item_widget = self._item_class(*args, **kwargs)
         item_widget.expansion_state_changed.connect(partial(_updatesize, item, item_widget))
         item_widget.imminent_deletion.connect(partial(_remove_item, item))
         item_widget.broadcast_size_hint.connect(item.setSizeHint)
 
-        self.setItemWidget(item, item_widget)
         self.addItem(item)
+        self.setItemWidget(item, item_widget)
         self._item_dict[id] = self.item_tuple(item, item_widget)
 
         item.setSizeHint(item_widget.sizeHint())
         return item_widget
 
+    def sort_by_key(self, role=_SORT_KEY_ROLE, order=Qt.SortOrder.AscendingOrder):
+        items = [self.takeItem(0) for i in range(self.count())]
+        items.sort(key=lambda it: it.data(role), reverse=(order == Qt.SortOrder.DescendingOrder))
+
+        for it in items:
+            self.addItem(it)
+            # reattach its custom widget
+            widget = self.itemWidget(it)
+            if widget:
+                self.setItemWidget(it, widget)
+
     def item_widget_pairs(self):
         return self._item_dict.values()
 

bec_widgets/widgets/control/device_manager/components/available_device_resources/available_device_group_ui.py

@@ -42,6 +42,7 @@ class Ui_AvailableDeviceGroup(object):
         title_layout.addWidget(self.add_all_button)
 
         self.device_list = QListWidget(AvailableDeviceGroup)
+        self.device_list.setSelectionMode(QListWidget.SelectionMode.ExtendedSelection)
         self.device_list.setObjectName("device_list")
         self.device_list.setFrameStyle(0)
 

bec_widgets/widgets/control/device_manager/components/available_device_resources/available_device_resources.py

@@ -22,38 +22,52 @@ class AvailableDeviceResources(BECWidget, QWidget, Ui_availableDeviceResources):
         super().__init__(parent=parent, **kwargs)
         self.setupUi(self)
         self._backend = get_backend()
-        self.refresh_full_list()
-        self.search_box.textChanged.connect(self.tag_groups_list.update_filter)
+        self.grouping_selector.addItem("deviceTags")
+        self.grouping_selector.addItems(self._backend.allowed_sort_keys)
+        self._grouping_selection_changed("deviceTags")
+        self.grouping_selector.currentTextChanged.connect(self._grouping_selection_changed)
+        self.search_box.textChanged.connect(self.device_groups_list.update_filter)
 
-    def refresh_full_list(self):
-        self.tag_groups_list.clear()
-        for tag_group, devices in self._backend.tag_groups.items():
-            self._add_tag_group(tag_group, devices)
-        self._add_tag_group("Untagged devices", self._backend.untagged_devices)
+    def refresh_full_list(self, device_groups: dict[str, set[HashableDevice]]):
+        self.device_groups_list.clear()
+        for device_group, devices in device_groups.items():
+            self._add_device_group(device_group, devices)
+        if self.grouping_selector.currentText == "deviceTags":
+            self._add_device_group("Untagged devices", self._backend.untagged_devices)
+        self.device_groups_list.sort_by_key()
 
-    def _add_tag_group(self, tag_group: str, devices: set[HashableDevice]):
-        self.tag_groups_list.add_item(
-            tag_group, self.tag_groups_list, tag_group, devices, expanded=False
+    def _add_device_group(self, device_group: str, devices: set[HashableDevice]):
+        self.device_groups_list.add_item(
+            device_group, self.device_groups_list, device_group, devices, expanded=False
         )
 
     def _reset_devices_state(self):
-        for tag_group in self.tag_groups_list.widgets():
-            tag_group.reset_devices_state()
+        for device_group in self.device_groups_list.widgets():
+            device_group.reset_devices_state()
 
     def set_devices_state(self, devices: Iterable[HashableDevice], included: bool):
         for device in devices:
-            for tag_group in self.tag_groups_list.widgets():
-                tag_group.set_item_state(hash(device), included)
+            for device_group in self.device_groups_list.widgets():
+                device_group.set_item_state(hash(device), included)
 
     def resizeEvent(self, event):
         super().resizeEvent(event)
-        for list_item, tag_group_widget in self.tag_groups_list.item_widget_pairs():
-            list_item.setSizeHint(tag_group_widget.sizeHint())
+        for list_item, device_group_widget in self.device_groups_list.item_widget_pairs():
+            list_item.setSizeHint(device_group_widget.sizeHint())
 
     @SafeSlot(list)
     def update_devices_state(self, config_list: list[dict[str, Any]]):
         self.set_devices_state(yield_only_passing(HashableDevice.model_validate, config_list), True)
 
+    @SafeSlot(str)
+    def _grouping_selection_changed(self, sort_key: str):
+        self.search_box.setText("")
+        if sort_key == "deviceTags":
+            device_groups = self._backend.tag_groups
+        else:
+            device_groups = self._backend.group_by_key(sort_key)
+        self.refresh_full_list(device_groups)
+
 
 if __name__ == "__main__":
     import sys

bec_widgets/widgets/control/device_manager/components/available_device_resources/available_device_resources_ui.py

@@ -30,24 +30,26 @@ class Ui_availableDeviceResources(object):
         self.search_layout.addWidget(self.search_box)
         self.search_layout.addWidget(QLabel("Group by: "))
         self.grouping_selector = QComboBox()
-        self.grouping_selector.addItems(["deviceTags", "deviceClass"])
+        self.search_layout.addWidget(self.grouping_selector)
 
-        self.tag_groups_list = ListOfExpandableFrames(
+        self.device_groups_list = ListOfExpandableFrames(
             availableDeviceResources, AvailableDeviceGroup
         )
-        self.tag_groups_list.setObjectName("tag_groups_list")
-        self.tag_groups_list.setSelectionMode(QAbstractItemView.SelectionMode.NoSelection)
-        self.tag_groups_list.setVerticalScrollMode(QAbstractItemView.ScrollMode.ScrollPerPixel)
-        self.tag_groups_list.setHorizontalScrollMode(QAbstractItemView.ScrollMode.ScrollPerPixel)
-        self.tag_groups_list.setMovement(QListView.Movement.Static)
-        self.tag_groups_list.setSpacing(2)
-        self.tag_groups_list.setDragDropMode(QListWidget.DragDropMode.DragOnly)
-        self.tag_groups_list.setDragEnabled(True)
-        self.tag_groups_list.setAcceptDrops(False)
-        self.tag_groups_list.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
+        self.device_groups_list.setObjectName("device_groups_list")
+        self.device_groups_list.setSelectionMode(QAbstractItemView.SelectionMode.NoSelection)
+        self.device_groups_list.setVerticalScrollMode(QAbstractItemView.ScrollMode.ScrollPerPixel)
+        self.device_groups_list.setHorizontalScrollMode(QAbstractItemView.ScrollMode.ScrollPerPixel)
+        self.device_groups_list.setMovement(QListView.Movement.Static)
+        self.device_groups_list.setSpacing(2)
+        self.device_groups_list.setDragDropMode(QListWidget.DragDropMode.DragOnly)
+        self.device_groups_list.setSelectionBehavior(QListWidget.SelectionBehavior.SelectItems)
+        self.device_groups_list.setSelectionMode(QListWidget.SelectionMode.ExtendedSelection)
+        self.device_groups_list.setDragEnabled(True)
+        self.device_groups_list.setAcceptDrops(False)
+        self.device_groups_list.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
         availableDeviceResources.setMinimumWidth(250)
         availableDeviceResources.resize(250, availableDeviceResources.height())
 
-        self.verticalLayout.addWidget(self.tag_groups_list)
+        self.verticalLayout.addWidget(self.device_groups_list)
 
         QMetaObject.connectSlotsByName(availableDeviceResources)

bec_widgets/widgets/control/device_manager/components/available_device_resources/device_resource_backend.py

@@ -42,6 +42,11 @@ class DeviceResourceBackend(Protocol):
         """A set of all untagged devices. The same device may not appear more than once."""
         ...
 
+    @property
+    def allowed_sort_keys(self) -> set[str]:
+        """A set of all fields which you may group devices by"""
+        ...
+
     def tags(self) -> set[str]:
         """Returns a set of all the tags in all available devices."""
         ...
@@ -50,6 +55,11 @@ class DeviceResourceBackend(Protocol):
         """Returns a set of the devices in the tag group with the given key."""
         ...
 
+    def group_by_key(self, key: str) -> dict[str, set[HashableDevice]]:
+        """Return a dict of all devices, organised by the specified key, which must be one of
+        the string keys in the Device model."""
+        ...
+
 
 def _devices_from_file(file: str, include_source: bool = True):
     data = yaml_load(file, process_includes=False)
@@ -68,7 +78,7 @@ class _ConfigFileBackend(DeviceResourceBackend):
         ] = self._get_config_from_backup_files() | self._get_configs_from_plugin_files(
             Path(plugin_repo_path()) / plugin_package_name() / "device_configs/"
         )
-        self._tag_groups = self._get_tag_groups()
+        self._device_groups = self._get_tag_groups()
 
     def _get_config_from_backup_files(self):
         dir = _BASE_REPO_PATH / "logs/device_configs/recovery_configs"
@@ -90,7 +100,7 @@ class _ConfigFileBackend(DeviceResourceBackend):
 
     @property
     def tag_groups(self):
-        return self._tag_groups
+        return self._device_groups
 
     @property
     def all_devices(self):
@@ -100,12 +110,22 @@ class _ConfigFileBackend(DeviceResourceBackend):
     def untagged_devices(self):
         return {d for d in self._raw_device_set if d.deviceTags == set()}
 
+    @property
+    def allowed_sort_keys(self) -> set[str]:
+        return {n for n, info in HashableDevice.model_fields.items() if info.annotation is str}
+
     def tags(self) -> set[str]:
         return reduce(operator.or_, (dev.deviceTags for dev in self._raw_device_set))
 
     def tag_group(self, tag: str) -> set[HashableDevice]:
         return self.tag_groups[tag]
 
+    def group_by_key(self, key: str) -> dict[str, set[HashableDevice]]:
+        if key not in self.allowed_sort_keys:
+            raise ValueError(f"Cannot group available devices by model key {key}")
+        group_names: set[str] = {getattr(item, key) for item in self._raw_device_set}
+        return {g: {d for d in self._raw_device_set if getattr(d, key) == g} for g in group_names}
+
 
 def get_backend() -> DeviceResourceBackend:
     return _ConfigFileBackend()