180 lines
6.8 KiB
Python
180 lines
6.8 KiB
Python
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
|
|
|
|
|
|
@dataclass
|
|
class DocumentTreeNode:
|
|
name: str
|
|
value: object
|
|
component_id: ComponentID | None
|
|
parent: DocumentTreeNode | None
|
|
children: list[DocumentTreeNode]
|
|
|
|
|
|
class DocumentTreeModel(QAbstractItemModel):
|
|
rename_document_requested = Signal(str)
|
|
rename_component_requested = Signal(Component, str)
|
|
|
|
def __init__(self, editable: bool = True) -> None:
|
|
super().__init__()
|
|
self._editable = editable
|
|
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.set_documents([document])
|
|
|
|
def set_documents(self, documents: list[CoreDocument]) -> None:
|
|
self.beginResetModel()
|
|
self._component_icons = {}
|
|
self._component_nodes = {}
|
|
self._root = self._build_tree(documents)
|
|
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 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 not in (0,1) 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 not isinstance(node, DocumentTreeNode):
|
|
return None
|
|
|
|
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() or index.column() != 0:
|
|
return False
|
|
|
|
node = index.internalPointer()
|
|
if not isinstance(node, DocumentTreeNode) or not isinstance(node.value, (CoreDocument, Component)):
|
|
return False
|
|
|
|
name = str(value).strip()
|
|
|
|
if not name or name == node.name:
|
|
return False
|
|
|
|
if isinstance(node.value, CoreDocument):
|
|
self.rename_document_requested.emit(name)
|
|
elif isinstance(node.value, Component):
|
|
self.rename_component_requested.emit(node.value, name)
|
|
return True
|
|
|
|
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 value(self, index: QModelIndex) -> object | None:
|
|
if not index.isValid():
|
|
return None
|
|
node = index.internalPointer()
|
|
return node.value if isinstance(node, DocumentTreeNode) else None
|
|
|
|
def component_index(self, component_id: ComponentID) -> QModelIndex:
|
|
node = self._component_nodes.get(component_id)
|
|
if node is None or node.parent is None:
|
|
return QModelIndex()
|
|
return self.createIndex(node.parent.children.index(node), 0, node)
|
|
|
|
def flags(self, index: QModelIndex) -> Qt.ItemFlag:
|
|
flags = super().flags(index)
|
|
|
|
if not index.isValid():
|
|
return flags
|
|
|
|
node = index.internalPointer()
|
|
|
|
# Make the document root node editable
|
|
if self._editable and index.column() == 0 and isinstance(node, DocumentTreeNode) and isinstance(node.value, (CoreDocument, Component)):
|
|
flags |= Qt.ItemFlag.ItemIsEditable
|
|
|
|
return flags
|
|
|
|
def _build_tree(self, documents: list[CoreDocument]) -> DocumentTreeNode:
|
|
# QT's invisible root
|
|
root = DocumentTreeNode(
|
|
name="",
|
|
value=None,
|
|
component_id=None,
|
|
parent=None,
|
|
children=[],
|
|
)
|
|
def _list_children(parent: DocumentTreeNode, components: dict[ComponentID, Component]) -> None:
|
|
for component_id, component in sorted(components.items(), key=lambda item: (item[1].name.casefold(), item[1].name, str(item[0]))):
|
|
component_node = DocumentTreeNode(name=component.name, value=component, component_id=component_id, parent=parent, children=[])
|
|
parent.children.append(component_node)
|
|
self._component_nodes[component_id] = component_node
|
|
|
|
if isinstance(component.implementation, GraphImplementation):
|
|
_list_children(component_node, component.implementation.graph.components)
|
|
|
|
for document in documents:
|
|
document_root = DocumentTreeNode(name=document.name, value=document, component_id=None, parent=root, children=[])
|
|
root.children.append(document_root)
|
|
_list_children(document_root, document.root)
|
|
|
|
return root
|