124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
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
|