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
@@ -34,6 +35,8 @@ class _TTLWidget(QtWidgets.QFrame):
         self.setFrameShape(QtWidgets.QFrame.Box)
         self.setFrameShadow(QtWidgets.QFrame.Raised)
 
+        self.uid = title
+
         grid = QtWidgets.QGridLayout()
         grid.setContentsMargins(2, 2, 2, 2)
         grid.setHorizontalSpacing(2)
@@ -112,6 +115,7 @@ class _TTLHandler:
         self.force_out = force_out
         self.set_mode = dm.ttl_set_mode
         self.title = title
+        self.uid = title
         self.cur_level = False
         self.cur_oe = False
         self.cur_override = False
@@ -165,6 +169,8 @@ class _DDSWidget(QtWidgets.QFrame):
     def __init__(self, title, is_urukul):
         QtWidgets.QFrame.__init__(self)
 
+        self.uid = title
+
         self.setFrameShape(QtWidgets.QFrame.Box)
         self.setFrameShadow(QtWidgets.QFrame.Raised)
 
@@ -282,6 +288,7 @@ class _DDSHandler:
                  ref_clk, cpld=None, pll=1, clk_div=0):
         self.dm = dm
         self.title = title
+        self.uid = title
         self.bus_channel = bus_channel
         self.channel = channel
         self.cur_frequency = 0
@@ -351,6 +358,8 @@ class _DACWidget(QtWidgets.QFrame):
     def __init__(self, channel, title):
         QtWidgets.QFrame.__init__(self)
 
+        self.uid = (title, channel)
+
         self.setFrameShape(QtWidgets.QFrame.Box)
         self.setFrameShadow(QtWidgets.QFrame.Raised)
 
@@ -378,6 +387,7 @@ class _DACHandler:
         self.spi_channel = spi_channel
         self.channel = channel
         self.title = title
+        self.uid = (title, channel)
 
     def create_widget(self):
         self.widget = _DACWidget(self.channel, self.title)
@@ -819,46 +829,78 @@ class _DeviceManager:
             await self.mi_connection.close()
 
 
-class MoninjTree(QtWidgets.QTreeWidget):
+class MoninjTreeWidget(pg.TreeWidget):
     def __init__(self):
-        QtWidgets.QTreeWidget.__init__(self)
+        pg.TreeWidget.__init__(self)
-        self.setDragDropMode(QtWidgets.QAbstractItemView.InternalMove)
+        self.setColumnCount(1)
-        self._groups = dict()
+        self.setSelectionMode(QtWidgets.QAbstractItemView.SingleSelection)
-        self._widgets = dict()
 
-    def add_widget(self, widget, group=None):
+        palette = QtGui.QPalette()
-        item = QtWidgets.QTreeWidgetItem([widget.to_model_path()])
+        palette.setColor(QtGui.QPalette.Highlight, QtGui.QColor("gray"))
-        self.setItemWidget(item, widget)
+        self.setPalette(palette)
-        self._widgets[widget.to_model_path()] = widget
+
-        if group is not None:
+        self._widgets = dict()
-            self._get_group(group).addChild(item)
+        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("Add new group", self)
+        add_grp_action.triggered.connect(self.add_group)
+        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 add_widget(self, widget):
+        item = pg.TreeWidgetItem()
+        item.setFlags(item.flags() & ~QtCore.Qt.ItemFlag.ItemIsDropEnabled)
+        self._widgets[widget.uid] = widget
+        self.addTopLevelItem(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:
+            self.remove_group_items(item)
         else:
-            self.addTopLevelItem(item)
+            self.remove_widget(item)
+        root = self.invisibleRootItem()
+        (item.parent() or root).removeChild(item)
 
     def remove_widget(self, item):
-        widget = self.itemWidget(child, 0)
+        widget = self.itemWidget(item, 0)
-        self.removeItemWidget(child, 0)
+        del self._widgets[widget.uid]
-        del self._widgets[widget.to_model_path()]
-        widget.deleteLater()
 
-    def add_group(self, group=None):
+    def add_group(self):
-        if group in self._groups:
+        group = "untitled_group"
-            raise SyntaxError("group name repeated")
+        group_item = pg.TreeWidgetItem([group])
-        group_item = QtWidgets.QTreeWidgetItem([group])
+        group_item.setFlags(group_item.flags() | QtCore.Qt.ItemFlag.ItemIsEditable)
         self.addTopLevelItem(group_item)
-        self._groups[group] = group_item
         return group_item
 
-    def remove_group(self, group):
+    def remove_group_items(self, group_item):
-        group_item = self._groups[group]
         child_items = group_item.takeChildren()
         for child in child_items:
             self.remove_widget(child)
-        self.takeTopLevelItem(self.indexOfTopLevelItem(group_item))
-        del self._groups[group]
-
-    def _get_group(self, group):
-        return self._groups[group]
 
 
 class MonInjDock(QtWidgets.QDockWidget):
@@ -884,21 +926,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.tree = MoninjTreeWidget()
-        self.grid = QtWidgets.QVBoxLayout()
+        layout.addWidget(self.tree, 1, 0, colspan=10)
-        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)
 
     def layout_widgets(self, handlers):
-        self.grid.removeItem(self.spacer)
         for handler in handlers:
-            self.grid.addWidget(handler.widget)
+            handler.create_widget()
-        self.grid.insertItem(-1, self.spacer)
+            self.tree.add_widget(handler.widget)
+            handler.refresh_display()
 
     def set_channels(self, handlers):
         self._channel_model.update(handlers)