Preserve TreeView state during keyed reorders

Author: berendkleinhaneveld

collagraph/renderers/pyside/objects/itemmodel.py

@@ -1,8 +1,76 @@
-from PySide6.QtGui import QStandardItemModel
+from PySide6.QtCore import QItemSelectionModel
+from PySide6.QtGui import QStandardItem, QStandardItemModel
+from PySide6.QtWidgets import QListView, QTableView, QTreeView
 
 from ... import PySideRenderer
 
 
+def _get_view_from_model(model: QStandardItemModel):
+    parent = model.parent()
+    if isinstance(parent, (QTreeView, QListView, QTableView)):
+        return parent
+    return None
+
+
+def _snapshot_item_state(
+    item: QStandardItem,
+    model: QStandardItemModel,
+    view,
+) -> list[tuple[QStandardItem, bool, bool]]:
+    snapshot: list[tuple[QStandardItem, bool, bool]] = []
+
+    def walk(it: QStandardItem) -> None:
+        index = model.indexFromItem(it)
+        selected = view.selectionModel().isSelected(index)
+        expanded = isinstance(view, QTreeView) and view.isExpanded(index)
+        snapshot.append((it, expanded, selected))
+        for row in range(it.rowCount()):
+            child = it.child(row)
+            if child:
+                walk(child)
+
+    walk(item)
+    return snapshot
+
+
+def _apply_item_selection(
+    model: QStandardItemModel,
+    view,
+    item: QStandardItem,
+    selected: bool,
+) -> None:
+    selection_model = view.selectionModel()
+    index = model.indexFromItem(item)
+    flags = QItemSelectionModel.SelectionFlag.Rows
+    if selected:
+        flags |= QItemSelectionModel.SelectionFlag.Select
+    else:
+        flags |= QItemSelectionModel.SelectionFlag.Deselect
+    selection_model.select(index, flags)
+
+
+def _apply_item_expanded(
+    model: QStandardItemModel,
+    view,
+    item: QStandardItem,
+    expanded: bool,
+) -> None:
+    if isinstance(view, QTreeView):
+        index = model.indexFromItem(item)
+        view.setExpanded(index, expanded)
+
+
+def _restore_item_state(
+    snapshot: list[tuple[QStandardItem, bool, bool]],
+    model: QStandardItemModel,
+    view,
+) -> None:
+    for item, expanded, _selected in snapshot:
+        _apply_item_expanded(model, view, item, expanded)
+    for item, _expanded, selected in snapshot:
+        _apply_item_selection(model, view, item, selected)
+
+
 @PySideRenderer.register_insert(QStandardItemModel)
 def insert(self, el, anchor=None):
     if isinstance(self, QStandardItemModel):
@@ -16,11 +84,29 @@ def insert(self, el, anchor=None):
             self.insertRow(index.row(), el)
         else:
             self.appendRow(el)
+
+        view = _get_view_from_model(self)
+        if view:
+            if hasattr(el, "_state_snapshot"):
+                _restore_item_state(el._state_snapshot, self, view)
+                del el._state_snapshot
+
+            if hasattr(el, "_expanded"):
+                _apply_item_expanded(self, view, el, el._expanded)
+                delattr(el, "_expanded")
+
+            if hasattr(el, "_selected"):
+                _apply_item_selection(self, view, el, el._selected)
+                delattr(el, "_selected")
     else:
         raise NotImplementedError(type(self).__name__)
 
 
 @PySideRenderer.register_remove(QStandardItemModel)
 def remove(self, el):
+    view = _get_view_from_model(self)
+    if view:
+        el._state_snapshot = _snapshot_item_state(el, self, view)
+
     index = self.indexFromItem(el)
     self.takeRow(index.row())

collagraph/renderers/pyside/objects/standarditem.py

@@ -1,9 +1,98 @@
+from PySide6.QtCore import QItemSelectionModel
 from PySide6.QtGui import QStandardItem
+from PySide6.QtWidgets import QListView, QTableView, QTreeView
 
 from ... import PySideRenderer
 from .qobject import set_attribute as qobject_set_attribute
 
 
+def _get_view_from_item(item: QStandardItem):
+    model = item.model()
+    if not model:
+        return None
+
+    parent = model.parent()
+    if isinstance(parent, (QTreeView, QListView, QTableView)):
+        return parent
+    return None
+
+
+def _snapshot_item_state(item: QStandardItem):
+    model = item.model()
+    view = _get_view_from_item(item)
+    if not (model and view):
+        return None
+
+    snapshot: list[tuple[QStandardItem, bool, bool]] = []
+
+    def walk(it: QStandardItem) -> None:
+        index = model.indexFromItem(it)
+        selected = view.selectionModel().isSelected(index)
+        expanded = isinstance(view, QTreeView) and view.isExpanded(index)
+        snapshot.append((it, expanded, selected))
+        for row in range(it.rowCount()):
+            child = it.child(row)
+            if child:
+                walk(child)
+
+    walk(item)
+    return snapshot
+
+
+def _apply_item_selection(item: QStandardItem, selected: bool) -> None:
+    model = item.model()
+    view = _get_view_from_item(item)
+    if not (model and view):
+        item._selected = selected
+        return
+
+    selection_model = view.selectionModel()
+    index = model.indexFromItem(item)
+    flags = QItemSelectionModel.SelectionFlag.Rows
+    if selected:
+        flags |= QItemSelectionModel.SelectionFlag.Select
+    else:
+        flags |= QItemSelectionModel.SelectionFlag.Deselect
+    selection_model.select(index, flags)
+
+
+def _apply_item_expanded(item: QStandardItem, expanded: bool) -> None:
+    model = item.model()
+    view = _get_view_from_item(item)
+    if not (model and isinstance(view, QTreeView)):
+        item._expanded = expanded
+        return
+
+    index = model.indexFromItem(item)
+    view.setExpanded(index, expanded)
+
+
+def _restore_item_state(item: QStandardItem) -> None:
+    if not hasattr(item, "_state_snapshot"):
+        return
+
+    model = item.model()
+    view = _get_view_from_item(item)
+    if not (model and view):
+        return
+
+    snapshot = item._state_snapshot
+    for it, expanded, _selected in snapshot:
+        if isinstance(view, QTreeView):
+            index = model.indexFromItem(it)
+            view.setExpanded(index, expanded)
+    for it, _expanded, selected in snapshot:
+        index = model.indexFromItem(it)
+        flags = QItemSelectionModel.SelectionFlag.Rows
+        if selected:
+            flags |= QItemSelectionModel.SelectionFlag.Select
+        else:
+            flags |= QItemSelectionModel.SelectionFlag.Deselect
+        view.selectionModel().select(index, flags)
+
+    del item._state_snapshot
+
+
 @PySideRenderer.register_insert(QStandardItem)
 def insert(self, el, anchor=None):
     if hasattr(el, "model_index"):
@@ -19,13 +108,29 @@ def insert(self, el, anchor=None):
                 break
         if index is None:
             return
+        if el.row() >= 0 and el.parent() == self:
+            self.takeRow(el.row())
         self.insertRow(index, el)
     else:
         self.appendRow(el)
 
+    _restore_item_state(el)
+
+    if hasattr(el, "_expanded"):
+        _apply_item_expanded(el, el._expanded)
+        delattr(el, "_expanded")
+
+    if hasattr(el, "_selected"):
+        _apply_item_selection(el, el._selected)
+        delattr(el, "_selected")
+
 
 @PySideRenderer.register_remove(QStandardItem)
 def remove(self, el):
+    snapshot = _snapshot_item_state(el)
+    if snapshot is not None:
+        el._state_snapshot = snapshot
+
     if hasattr(el, "model_index"):
         # Only support removal of rows for now
         row, _column = getattr(el, "model_index")
@@ -48,6 +153,14 @@ def set_attribute(self, attr, value):
         model.blockSignals(True)
 
     try:
+        match attr:
+            case "selected":
+                _apply_item_selection(self, value)
+                return
+            case "expanded":
+                _apply_item_expanded(self, value)
+                return
+
         if attr == "model_index" and model:
             index = model.indexFromItem(self)
             row, column = value

examples/pyside/keyed_tree_demo.cgx

@@ -0,0 +1,135 @@
+<!--
+  Run with:
+  uv run collagraph examples/pyside/keyed_tree_demo.cgx
+
+  Try this sequence:
+  1) Expand a couple of parents and select a parent/child in the tree.
+  2) Click Reverse or Shuffle.
+  3) Observe that expanded + selection state stays with the moved items.
+-->
+<widget>
+  <label text="Keyed TreeView Reorder Demo" />
+  <label text="Manually expand/select items, then reorder parents." />
+
+  <widget :layout="{'type': 'Box', 'direction': 'LeftToRight'}">
+    <button
+      text="Reverse Parents"
+      @clicked="reverse_parents"
+    />
+    <button
+      text="Shuffle Parents"
+      @clicked="shuffle_parents"
+    />
+    <button
+      text="Print UI State"
+      @clicked="print_ui_state"
+    />
+  </widget>
+
+  <treeview object-name="tree">
+    <itemmodel>
+      <standarditem
+        v-for="parent in items"
+        :key="parent['id']"
+        :text="parent['text']"
+      >
+        <standarditem
+          v-for="child in parent['children']"
+          :key="child['id']"
+          :text="child['text']"
+        />
+      </standarditem>
+    </itemmodel>
+  </treeview>
+
+  <label :text="status_text()" />
+</widget>
+
+<script>
+import random
+
+import collagraph as cg
+from PySide6 import QtWidgets
+
+
+class App(cg.Component):
+    def init(self):
+        self.state["items"] = [
+            {
+                "id": 1,
+                "text": "Parent 1",
+                "children": [
+                    {"id": 11, "text": "Child 1.1"},
+                    {"id": 12, "text": "Child 1.2"},
+                ],
+            },
+            {
+                "id": 2,
+                "text": "Parent 2",
+                "children": [
+                    {"id": 21, "text": "Child 2.1"},
+                    {"id": 22, "text": "Child 2.2"},
+                ],
+            },
+            {
+                "id": 3,
+                "text": "Parent 3",
+                "children": [
+                    {"id": 31, "text": "Child 3.1"},
+                ],
+            },
+            {
+                "id": 4,
+                "text": "Parent 4",
+                "children": [
+                    {"id": 41, "text": "Child 4.1"},
+                    {"id": 42, "text": "Child 4.2"},
+                    {"id": 43, "text": "Child 4.3"},
+                ],
+            },
+        ]
+
+    def reverse_parents(self):
+        self.state["items"] = list(reversed(self.state["items"]))
+
+    def shuffle_parents(self):
+        items = list(self.state["items"])
+        random.shuffle(items)
+        self.state["items"] = items
+
+    def status_text(self):
+        parent_count = len(self.state["items"])
+        child_count = sum(len(parent["children"]) for parent in self.state["items"])
+        return f"Parents: {parent_count} | Children: {child_count}"
+
+    def print_ui_state(self):
+        tree = self.element.findChild(QtWidgets.QTreeView, "tree")
+        if not tree:
+            print("TreeView not found")
+            return
+
+        model = tree.model()
+        selection_model = tree.selectionModel()
+
+        selected = []
+        expanded = []
+
+        for row in range(model.rowCount()):
+            parent_item = model.item(row)
+            parent_index = model.indexFromItem(parent_item)
+
+            if selection_model.isSelected(parent_index):
+                selected.append(parent_item.text())
+            if tree.isExpanded(parent_index):
+                expanded.append(parent_item.text())
+
+            for child_row in range(parent_item.rowCount()):
+                child_item = parent_item.child(child_row)
+                child_index = model.indexFromItem(child_item)
+                if selection_model.isSelected(child_index):
+                    selected.append(child_item.text())
+
+        print(f"Order: {[parent['text'] for parent in self.state['items']]}")
+        print(f"Selected: {selected}")
+        print(f"Expanded parents: {expanded}")
+</script>