From 8b105c064a811139361c125852119724262fd3d3 Mon Sep 17 00:00:00 2001
From: Simon Renblad <srenblad@m-labs.hk>
Date: Wed, 17 Apr 2024 17:04:26 +0800
Subject: [PATCH] moninj: add MoninjTreeWidget

---
 artiq/dashboard/moninj.py | 95 +++++++++++++++++++++++++++++++++------
 1 file changed, 82 insertions(+), 13 deletions(-)

diff --git a/artiq/dashboard/moninj.py b/artiq/dashboard/moninj.py
index 9674ddb98..f5146953e 100644
--- a/artiq/dashboard/moninj.py
+++ b/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
@@ -828,6 +829,80 @@ 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)
+
+        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("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.remove_widget(item)
+        root = self.invisibleRootItem()
+        (item.parent() or root).removeChild(item)
+
+    def remove_widget(self, item):
+        widget = self.itemWidget(item, 0)
+        del self._widgets[widget.uid]
+
+    def add_group(self):
+        group = "untitled_group"
+        group_item = pg.TreeWidgetItem([group])
+        group_item.setFlags(group_item.flags() | QtCore.Qt.ItemFlag.ItemIsEditable)
+        self.addTopLevelItem(group_item)
+        return group_item
+
+    def remove_group_items(self, group_item):
+        child_items = group_item.takeChildren()
+        for child in child_items:
+            self.remove_widget(child)
+
+
 class MonInjDock(QtWidgets.QDockWidget):
     def __init__(self, schedule_ctl):
         QtWidgets.QDockWidget.__init__(self, "MonInj")
@@ -851,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.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.update(handlers)