QTreeView QAbstractItemModel CustomModel展示 pandas数据

# 参考
# https://gist.github.com/nbassler/342fc56c42df27239fa5276b79fca8e6
# https://blog.csdn.net/zyhse/article/details/105894431

from PyQt5 import QtCore, QtWidgets
import sys
import pandas as pd
import time

# 适应 二级tree .其他不适应
# 输入格式固定
class CustomModel(QtCore.QAbstractItemModel):
    def __init__(self, root):
        QtCore.QAbstractItemModel.__init__(self)
        self._root = []
        self._data = []
        for k in root.keys():
            self._root.append(k)
            self._data.append(root[k])
        self._clcnt = self._data[0].shape[1]

        #id 标识唯一
        self.lev1Ids = []
        self.lev2Ids = []
        for v in self._root:
            self.lev1Ids.append(id(v))
        for v in self._data:
            self.lev2Ids.append(id(v))

    def rowCount(self, index):
        v = index.internalPointer()
        if index.isValid():
            if index.column() > 0:
                return 0
            if id(v) == id(self._root):
                return len(self._root)
            elif id(v) in self.lev1Ids:
                i = len(self._data[self._root.index(v)])
                return i
            elif id(v) in self.lev2Ids:
                return 0
                return v.shape[0]
        return len(self._root)


    def index(self, row, column, _parent=None):
        child = None
        if not _parent or not _parent.isValid():
            parent = self._root
            child = self._root[row]
        else:
            parent = _parent.internalPointer()
            if id(parent) == id(self._root):
                child = self._root[row]
            elif id(parent) in self.lev1Ids:
                child = self._data[self._root.index(parent)]
            elif id(parent) in self.lev2Ids:
                return QtCore.QModelIndex()
            else:
                return QtCore.QModelIndex()
        return QtCore.QAbstractItemModel.createIndex(self, row, column, child)

    def parent(self, index):
        p = None
        rows = None
        v = index.internalPointer()
        if index.isValid() and v is not None:
            if id(v) == id(self._root):
                return QtCore.QModelIndex()
            elif id(v) in self.lev1Ids:
                p = self._root
                rows = 0
            elif id(v) in self.lev2Ids:
                p = self._root[self.lev2Ids.index(id(v))]
                rows = len(self._root)
            if p is not None:
                return QtCore.QAbstractItemModel.createIndex(self, rows, 0, p)
        return QtCore.QModelIndex()


    def columnCount(self, index):
        v = index.internalPointer()
        if index.isValid():
            if id(v) == id(self._root):
                return 1
            elif id(v) in self.lev1Ids:
                return self._clcnt
            elif id(v) in self.lev2Ids:
                return self._clcnt
        return self._clcnt



    def data(self, index, role):
        if not index.isValid():
            return None
        v = index.internalPointer()
        if role == QtCore.Qt.DisplayRole:
            if v is None:
                return "N"
            if id(v) == id(self._root):
                return 'root'
            elif id(v) in self.lev1Ids:
                if index.column() > 0:
                    return None
                return self._root[self._root.index(v)]
            elif id(v) in self.lev2Ids:
                return v.iloc[index.row(),index.column()]
        return None

def gen_datas():
    items = []
    items.append(pd.DataFrame([
        ['11', '12'],
        ['14', '15']
    ]))
    items.append(pd.DataFrame([
        ['21', '22'],
        ['24', '25'],
        ['26', '27'],
        ['28', '29']
    ]))
    items.append(pd.DataFrame([
        ['31', '32']
    ]))
    # dic
    root = dict()
    root['1'] = items[0]
    root['2'] = items[1]
    root['3'] = items[2]
    return root

def gen_Big_datas():

    items = []
    for i in range(100):
        item = []
        for j in range(1000):
            item.append(["c1-{}-{}".format(i, j), "c2-{}-{}".format(i, j)])
        items.append(item)
    root = dict()
    for i in range(len(items)):
        root[str(i)] = pd.DataFrame(items[i])
    return root


class MyTree():
    def __init__(self):
        self.items = []
        self.items.append(pd.DataFrame([
            ['11','12'],
            ['14','15']
        ]))
        self.items.append(pd.DataFrame([
            ['21', '22'],
            ['24', '25'],
            ['26', '27'],
            ['28', '29']
        ]))
        self.items.append(pd.DataFrame([
            ['31', '32']
        ]))
        # dic
        self.root = dict()
        self.root['1'] = self.items[0]
        self.root['2'] = self.items[1]
        self.root['3'] = self.items[2]

        self.tw = QtWidgets.QTreeView()
        self.tw.setModel(CustomModel(self.root))

def test():
    app = QtWidgets.QApplication(sys.argv)
    mytree = MyTree()
    mytree.tw.show()
    sys.exit(app.exec_())

def test2():
    app = QtWidgets.QApplication(sys.argv)

    # root = gen_datas()
    root = gen_Big_datas()
    begin = time.time()
    tw = QtWidgets.QTreeView()
    tw.setModel(CustomModel(root))
    tw.show()
    print(time.time()-begin)
    sys.exit(app.exec_())

def test3():
    app = QtWidgets.QApplication(sys.argv)

    # root = gen_datas()
    root = gen_Big_datas()
    begin = time.time()
    treeWidget = QtWidgets.QTreeWidget()
    treeWidget.setColumnCount(3)
    for index, root_text in enumerate(root.keys()):
        rootw = QtWidgets.QTreeWidgetItem(treeWidget)
        rootw.setText(0, root_text)
        for i,child_datas in root[root_text].iterrows():
            child = QtWidgets.QTreeWidgetItem(rootw)
            for col, val1col in enumerate(child_datas):
                child.setText(1 + col, str(val1col))

    treeWidget.show()
    print(time.time() - begin)
    sys.exit(app.exec_())

if __name__ == "__main__":
    # test()
    test2()
    # test3()