Icon in document tree

src/bedit_gui/views/models/document_tree_model.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 
 from PySide6.QtCore import QAbstractItemModel, QModelIndex, Qt, Signal
+from PySide6.QtGui import QIcon
 
 from bedit_core.models import Component, ComponentID, GraphImplementation
 from bedit_core.models import Document as CoreDocument
@@ -12,6 +13,7 @@ from bedit_core.models import Document as CoreDocument
 class DocumentTreeNode:
     name: str
     value: object
+    component_id: ComponentID | None
     parent: DocumentTreeNode | None
     children: list[DocumentTreeNode]
 
@@ -23,23 +25,40 @@ class DocumentTreeModel(QAbstractItemModel):
     def __init__(self) -> None:
         super().__init__()
         self._document: CoreDocument | None = None
-        self._root = DocumentTreeNode("Document", None, None, [])
+        self._root = DocumentTreeNode("Document", None, None, None, [])
+        self._component_icons: dict[ComponentID, QIcon] = {}
+        self._component_nodes: dict[ComponentID, DocumentTreeNode] = {}
 
     def set_document(self, document: CoreDocument) -> None:
         self.beginResetModel()
         self._document = document
+        self._component_icons = {}
+        self._component_nodes = {}
         self._root = self._build_tree(document)
         self.endResetModel()
 
+    def set_component_icon(self, component_id: ComponentID, icon: QIcon) -> None:
+        node = self._component_nodes.get(component_id)
+        if node is None or node.parent is None:
+            return
+        self._component_icons[component_id] = icon
+        row = node.parent.children.index(node)
+        index = self.createIndex(row, 1, node)
+        self.dataChanged.emit(index, index, [Qt.ItemDataRole.DecorationRole])
+
     def rowCount(self, parent: QModelIndex | None = None) -> int:
+        if parent is not None and parent.isValid() and parent.column() != 0:
+            return 0
         return len(self._node(parent).children)
 
     def columnCount(self, _parent: QModelIndex | None = None) -> int:
-        return 1
+        return 2
 
     def index(self, row: int, column: int, parent: QModelIndex | None = None) -> QModelIndex:
+        if parent is not None and parent.isValid() and parent.column() != 0:
+            return QModelIndex()
         parent_node = self._node(parent)
-        if column != 0 or row < 0 or row >= len(parent_node.children):
+        if column not in (0,1) or row < 0 or row >= len(parent_node.children):
             return QModelIndex()
         return self.createIndex(row, column, parent_node.children[row])
 
@@ -69,13 +88,19 @@ class DocumentTreeModel(QAbstractItemModel):
         if not isinstance(node, DocumentTreeNode):
             return None
 
-        if role in (Qt.ItemDataRole.DisplayRole, Qt.ItemDataRole.EditRole):
-            return node.name
+        if index.column() == 0:
+            if role in (Qt.ItemDataRole.DisplayRole, Qt.ItemDataRole.EditRole):
+                return node.name
+        elif index.column() == 1:
+            if role == Qt.ItemDataRole.DecorationRole:
+                return self._component_icons.get(node.component_id)
+            if role == Qt.ItemDataRole.TextAlignmentRole:
+                return Qt.AlignmentFlag.AlignCenter
 
         return None
 
     def setData(self, index: QModelIndex, value: object, role: int = Qt.ItemDataRole.EditRole) -> bool:
-        if role != Qt.ItemDataRole.EditRole or not index.isValid():
+        if role != Qt.ItemDataRole.EditRole or not index.isValid() or index.column() != 0:
             return False
 
         node = index.internalPointer()
@@ -115,7 +140,7 @@ class DocumentTreeModel(QAbstractItemModel):
         node = index.internalPointer()
 
         # Make the document root node editable
-        if isinstance(node, DocumentTreeNode) and isinstance(node.value, (CoreDocument, Component)):
+        if index.column() == 0 and isinstance(node, DocumentTreeNode) and isinstance(node.value, (CoreDocument, Component)):
             flags |= Qt.ItemFlag.ItemIsEditable
 
         return flags
@@ -125,17 +150,19 @@ class DocumentTreeModel(QAbstractItemModel):
         root = DocumentTreeNode(
             name="",
             value=None,
+            component_id=None,
             parent=None,
             children=[],
         )
         # Add itself as a child so the document root is visible in the tree
-        document_root = DocumentTreeNode(name=document.name, value=document, parent=root, children=[])
+        document_root = DocumentTreeNode(name=document.name, value=document, component_id=None, parent=root, children=[])
         root.children.append(document_root)
 
         def _list_children(root: DocumentTreeNode, components: dict[ComponentID, Component]) -> None:
-            for component in components.values():
-                component_node = DocumentTreeNode(name=component.name, value=component, parent=root, children=[])
+            for component_id, component in components.items():
+                component_node = DocumentTreeNode(name=component.name, value=component, component_id=component_id, parent=root, children=[])
                 root.children.append(component_node)
+                self._component_nodes[component_id] = component_node
 
                 if isinstance(component.implementation, GraphImplementation):
                     _list_children(component_node, component.implementation.graph.components)