moninj: add MoninjTreeWidget

artiq/dashboard/moninj.py

@@ -3,7 +3,8 @@ import logging
 import textwrap
 from collections import namedtuple
 
-from PyQt5 import QtCore, QtWidgets
+from PyQt5 import QtCore, QtWidgets, QtGui
+import pyqtgraph as pg
 
 from artiq.coredevice.comm_moninj import CommMonInj, TTLOverride, TTLProbe
 from artiq.coredevice.ad9912_reg import AD9912_SER_CONF
@@ -794,6 +795,129 @@ class _DeviceManager:
             await self.mi_connection.close()
 
 
+class MoninjTreeWidget(pg.TreeWidget):
+    def __init__(self):
+        pg.TreeWidget.__init__(self)
+        self.setColumnCount(1)
+        self.setSelectionMode(QtWidgets.QAbstractItemView.SingleSelection)
+        self.header().setVisible(False)
+        self.setHorizontalScrollMode(self.ScrollPerPixel)
+        self.setVerticalScrollMode(self.ScrollPerPixel)
+
+        palette = QtGui.QPalette()
+        palette.setColor(QtGui.QPalette.Highlight, QtGui.QColor("gray"))
+        self.setPalette(palette)
+
+        self._widgets = dict()
+        self.setContextMenuPolicy(QtCore.Qt.ActionsContextMenu)
+        delete_action = QtWidgets.QAction("Delete", self)
+        delete_action.triggered.connect(self._remove_item_action)
+        delete_action.setShortcut("DEL")
+        delete_action.setShortcutContext(QtCore.Qt.WidgetShortcut)
+        self.addAction(delete_action)
+        add_grp_action = QtWidgets.QAction("Insert group", self)
+        add_grp_action.triggered.connect(self._add_group_action)
+        add_grp_action.setShortcut("CTRL+SHIFT+N")
+        add_grp_action.setShortcutContext(QtCore.Qt.WidgetShortcut)
+        self.addAction(add_grp_action)
+
+    def supportedDropActions(self):
+        # Moving items breaks widgets in pg.TreeWidget, perform only copies
+        return QtCore.Qt.CopyAction
+
+    def itemMoving(self, item, parent, index):
+        if parent is None:
+            return True
+        if self.indexOfTopLevelItem(parent) == -1:
+            return False
+        if self.itemWidget(item, 0) is None:
+            return False
+        if self.itemWidget(parent, 0) is not None:
+            return False
+        return True
+
+    def _find_group(self, group):
+        for i in range(self.topLevelItemCount()):
+            item = self.topLevelItem(i)
+            if item.text(0) == group:
+                return item
+        return self.invisibleRootItem()
+
+    def add_widget(self, widget, group=None):
+        item = pg.TreeWidgetItem()
+        item.setFlags(item.flags() & ~QtCore.Qt.ItemFlag.ItemIsDropEnabled)
+        if group is None:
+            self.addTopLevelItem(item)
+        else:
+            group_item = self._find_group(group)
+            group_item.addChild(item)
+        self.setItemWidget(item, 0, widget)
+
+    def _remove_item_action(self):
+        if len(self.selectedItems()) == 0:
+            return
+        item = self.selectedItems()[0]
+        self.remove_item(item)
+
+    def remove_item(self, item):
+        if self.itemWidget(item, 0) is None:
+            item.takeChildren()
+        root = self.invisibleRootItem()
+        (item.parent() or root).removeChild(item)
+
+    def _add_group_action(self):
+        if len(self.selectedItems()) == 0:
+            self.add_group()
+            return
+        item = self.selectedItems()[0]
+        index = self.indexOfTopLevelItem(item)
+        if index != -1:
+            self.insert_group(index + 1)
+        else:
+            parent = item.parent()
+            index = self.indexOfTopLevelItem(parent)
+            self.insert_group(index + 1)
+
+    def add_group(self, name="untitled"):
+        group_item = pg.TreeWidgetItem([name])
+        group_item.setFlags(group_item.flags() | QtCore.Qt.ItemFlag.ItemIsEditable)
+        self.addTopLevelItem(group_item)
+        return group_item
+
+    def insert_group(self, index, name="untitled"):
+        group_item = pg.TreeWidgetItem([name])
+        group_item.setFlags(group_item.flags() | QtCore.Qt.ItemFlag.ItemIsEditable)
+        self.insertTopLevelItem(index, group_item)
+        return group_item
+
+    def extend(self, d):
+        for item in d:
+            if isinstance(item, dict):
+                self.add_group(item["name"])
+                for widget in item["widgets"]:
+                    self.add_widget(widget, group=item["name"])
+            else:
+                self.add_widget(item)
+
+    def export_list(self):
+        export_list = list()
+        for i in range(self.topLevelItemCount()):
+            item = self.topLevelItem(i)
+            widget = self.itemWidget(item, 0)
+            if widget is None:
+                group = dict()
+                group["name"] = item.text(0)
+                group["widgets"] = list()
+                for j in range(item.childCount()):
+                    child = item.child(j)
+                    widget = self.itemWidget(child, 0)
+                    group["widgets"].append(widget.uid)
+                export_list.append(group)
+            else:
+                export_list.append(widget.uid)
+        return export_list
+
+
 class MonInjDock(QtWidgets.QDockWidget):
     def __init__(self, schedule_ctl):
         QtWidgets.QDockWidget.__init__(self, "MonInj")
@@ -817,21 +941,15 @@ class MonInjDock(QtWidgets.QDockWidget):
                 QtWidgets.QStyle.SP_FileDialogListView))
         add_channel_btn.clicked.connect(self.add_channel_dialog.open)
         layout.addWidget(add_channel_btn, 0, 0, colspan=1)
-        scroll_area = QtWidgets.QScrollArea()
-        layout.addWidget(scroll_area, 1, 0, colspan=10)
-        self.grid = QtWidgets.QVBoxLayout()
-        self.spacer = QtWidgets.QSpacerItem(50, 500, QtWidgets.QSizePolicy.Ignored,
-                                            QtWidgets.QSizePolicy.Fixed)
-        grid_widget = QtWidgets.QWidget()
-        grid_widget.setLayout(self.grid)
-        scroll_area.setWidgetResizable(True)
-        scroll_area.setWidget(grid_widget)
+
+        self.tree = MoninjTreeWidget()
+        layout.addWidget(self.tree, 1, 0, colspan=10)
 
     def layout_widgets(self, handlers):
-        self.grid.removeItem(self.spacer)
         for handler in handlers:
-            self.grid.addWidget(handler.widget)
-        self.grid.insertItem(-1, self.spacer)
+            handler.create_widget()
+            self.tree.add_widget(handler.widget)
+            handler.refresh_display()
 
     def set_channels(self, handlers):
         self._channel_model.clear()