trove/trovedb/qgui/trove_tree_model.py

from typing import Optional
from PySide6.QtCore import QAbstractItemModel, QModelIndex, Qt, QObject

from trovedb.trove import Note, TreeNote, Trove, ObjectId, TreeEntry


class TroveNode:
    """Wrapper around a Note, caching children for tree model use."""

    def __init__(self, note: Note, name: str, parent: 'Optional[TroveNode]' = None):
        self.note = note
        self.name = name
        self.parent = parent
        self._children: Optional[list[TroveNode]] = None  # None = not yet loaded

    def is_tree(self) -> bool:
        return isinstance(self.note, TreeNote)

    def children(self) -> 'list[TroveNode]':
        """Lazy-load and cache children. Sorted by name."""
        if self._children is None:
            if self.is_tree():
                tree: TreeNote = self.note  # type: ignore[assignment]
                entries: list[TreeEntry] = sorted(tree.entries(), key=lambda e: e.name)
                self._children = [
                    TroveNode(tree.child(e.name), e.name, parent=self)
                    for e in entries
                ]
            else:
                self._children = []
        return self._children

    def invalidate(self) -> None:
        """Clear cached children, forcing a reload on next access."""
        self._children = None

    def child_at(self, row: int) -> 'Optional[TroveNode]':
        kids = self.children()
        return kids[row] if 0 <= row < len(kids) else None

    def row(self) -> int:
        if self.parent is None:
            return 0
        return self.parent.children().index(self)

    def child_count(self) -> int:
        return len(self.children())


class TroveTreeModel(QAbstractItemModel):
    """
    Read-only tree model backed by a Trove instance.

    The root TreeNote is treated as an invisible root; its children
    are the top-level items in the view.
    """

    def __init__(self, trove: Trove, parent: Optional[QObject] = None):
        super().__init__(parent)
        root_note = trove.get_root()
        self._root = TroveNode(root_note, "/", parent=None)

    # ------------------------------------------------------------------
    # QAbstractItemModel interface
    # ------------------------------------------------------------------

    def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> QModelIndex:
        if not self.hasIndex(row, column, parent):
            return QModelIndex()

        parent_node = self._node_from_index(parent)
        child = parent_node.child_at(row)
        if child is None:
            return QModelIndex()
        return self.createIndex(row, column, child)

    def parent(self, index: QModelIndex) -> QModelIndex:  # type: ignore[override]
        if not index.isValid():
            return QModelIndex()

        node: TroveNode = index.internalPointer()  # type: ignore[assignment]
        parent_node = node.parent

        if parent_node is None or parent_node is self._root:
            return QModelIndex()

        return self.createIndex(parent_node.row(), 0, parent_node)

    def rowCount(self, parent: QModelIndex = QModelIndex()) -> int:
        parent_node = self._node_from_index(parent)
        return parent_node.child_count()

    def columnCount(self, parent: QModelIndex = QModelIndex()) -> int:
        return 1

    def data(self, index: QModelIndex, role: int = Qt.DisplayRole) -> object:  # type: ignore[assignment]
        if not index.isValid():
            return None
        node: TroveNode = index.internalPointer()  # type: ignore[assignment]
        if role == Qt.DisplayRole:
            return node.name
        if role == Qt.UserRole:
            return node.note
        return None

    def hasChildren(self, parent: QModelIndex = QModelIndex()) -> bool:
        node = self._node_from_index(parent)
        return node.is_tree()

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _node_from_index(self, index: QModelIndex) -> TroveNode:
        if index.isValid():
            node = index.internalPointer()
            if node is not None:
                return node  # type: ignore[return-value]
        return self._root

    def invalidate_node(self, index: QModelIndex) -> None:
        """
        Force a reload of the node at index.
        Call this after mutations to the underlying Trove tree.
        """
        node = self._node_from_index(index)
        self.beginResetModel()
        node.invalidate()
        self.endResetModel()