dictionary_item_model.py

import sys
import time
from typing import Dict, Optional, Tuple, Any

import numpy as np
from PySide2.QtCore import QAbstractItemModel, QModelIndex, Qt
from PySide2.QtWidgets import QApplication, QTreeView


class TreeItem:

    def __init__(self, path: Tuple, dict_data: Dict, parent_item: Optional["TreeItem"]):
        self.path = path
        self.key = path[-1] if len(path) > 0 else ''
        self.parent_item = parent_item
        self.child_items = {}
        self.dict_data = dict_data

    def child(self, row: int) -> "TreeItem":
        key = list(self.dict_data)[row]
        return self.child_items.setdefault(row, TreeItem((*self.path, key), self.dict_data[key], self))

    def row(self) -> int:
        return next(row for row, child in self.parent_item.child_items.items()
                    if child == self) if self.parent_item else 0

    def row_count(self) -> int:
        return len(self.dict_data) if isinstance(self.dict_data, Dict) else 0


class DictionaryItemModel(QAbstractItemModel):

    def __init__(self, dict_data: Dict, parent):
        super().__init__(parent)
        self.root_item = TreeItem((), dict_data, None)
        self.parents = {0: self.root_item}

    def rowCount(self, parent: QModelIndex = ...) -> int:
        if parent.column() > 0:
            return 0
        if not parent.isValid():
            parent_item = self.root_item
        else:
            parent_item = parent.internalPointer()
        return parent_item.row_count()

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

    def data(self, index: QModelIndex, role: int = ...) -> Any:
        if not index.isValid():
            return None
        item = index.internalPointer()

        if role == Qt.DisplayRole:
            if item.key is not None and index.column() == 0:
                return f'{item.key}:{item.path}'

        elif role == Qt.UserRole:
            if item:
                return item.key

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

        if not parent.isValid():
            parent_item = self.root_item
        else:
            parent_item = parent.internalPointer()

        child_item = parent_item.child(row)
        if child_item:
            return self.createIndex(row, column, child_item)
        else:
            return QModelIndex()

    def parent(self, index: QModelIndex) -> QModelIndex:
        if not index.isValid():
            return QModelIndex()

        child_item = index.internalPointer()
        if not child_item:
            return QModelIndex()

        parent_item = child_item.parent_item

        if parent_item == self.root_item:
            return QModelIndex()

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


def main():
    data = {
        str(i): {str(j): {str(k): np.zeros((1200, 1)) for k in range(50)} for j in range(50)} for i in range(50)
    }

    app = QApplication(sys.argv)
    start = time.time()
    tree_view = QTreeView()
    tree_model = DictionaryItemModel(data, tree_view)
    tree_view.setModel(tree_model)
    print(f'Loading model: {time.time() - start}')
    tree_view.show()
    app.exec_()


if __name__ == "__main__":
    main()