Document tree panel

src/bedit_gui/views/models/document_tree_model.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from PySide6.QtCore import QAbstractItemModel, QModelIndex, Qt
+
+from bedit_core.models import Document as CoreDocument, ComponentID, Component, GraphImplementation
+
+
+@dataclass
+class DocumentTreeNode:
+    name: str
+    value: object
+    parent: DocumentTreeNode | None
+    children: list[DocumentTreeNode]
+
+
+class DocumentTreeModel(QAbstractItemModel):
+    def __init__(self) -> None:
+        super().__init__()
+        self._document: CoreDocument | None = None
+        self._root = DocumentTreeNode("Document", None, None, [])
+
+    def set_document(self, document: CoreDocument) -> None:
+        self.beginResetModel()
+        self._document = document
+        self._root = self._build_tree(document)
+        self.endResetModel()
+
+    def rowCount(self, parent: QModelIndex | None = None) -> int:
+        return len(self._node(parent).children)
+
+    def columnCount(self, _parent: QModelIndex | None = None) -> int:
+        return 1
+
+    def index(
+        self,
+        row: int,
+        column: int,
+        parent: QModelIndex | None = None,
+    ) -> QModelIndex:
+        parent_node = self._node(parent)
+        if (
+            column != 0
+            or row < 0
+            or row >= len(parent_node.children)
+        ):
+            return QModelIndex()
+        return self.createIndex(row, column, parent_node.children[row])
+
+    def parent(self, index: QModelIndex) -> QModelIndex:
+        if not index.isValid():
+            return QModelIndex()
+
+        node = index.internalPointer()
+        if not isinstance(node, DocumentTreeNode):
+            return QModelIndex()
+
+        parent_node = node.parent
+        if parent_node is None or parent_node is self._root:
+            return QModelIndex()
+
+        grandparent = parent_node.parent
+        if grandparent is None:
+            return QModelIndex()
+        row = grandparent.children.index(parent_node)
+        return self.createIndex(row, 0, parent_node)
+
+    def data(
+        self,
+        index: QModelIndex,
+        role: int = Qt.ItemDataRole.DisplayRole,
+    ) -> object | None:
+        if not index.isValid():
+            return None
+
+        node = index.internalPointer()
+        if (
+            role == Qt.ItemDataRole.DisplayRole
+            and isinstance(node, DocumentTreeNode)
+        ):
+            return node.name
+        return None
+
+    def _node(self, index: QModelIndex | None) -> DocumentTreeNode:
+        if index is None or not index.isValid():
+            return self._root
+
+        node = index.internalPointer()
+        return node if isinstance(node, DocumentTreeNode) else self._root
+
+    def _build_tree(self, document: CoreDocument) -> DocumentTreeNode:
+        # QT's invisible root
+        root = DocumentTreeNode(
+            name="",
+            value=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 = []
+        )
+        root.children.append(document_root)
+
+        def _list_children(root: DocumentTreeNode, components: dict[ComponentID, Component]) -> None:
+            for component_id, component in components.items():
+                component_node = DocumentTreeNode(
+                    name = component.name,
+                    value = (component_id, component),
+                    parent = root,
+                    children = []
+                )
+                root.children.append(component_node)
+
+                if isinstance(component.implementation, GraphImplementation):
+                    _list_children(component_node, component.implementation.graph.components)
+        _list_children(document_root, document.root)
+
+        return root