pyqt5 qtreeview 结合自定义model添加按钮

 1. 序言

        在使用PyQt5开发的过程中，涉及到一个场景，需要在树中添加菜单按钮。经过研究学习，采用如下方案可以完成树中添加菜单按钮的功能，因过程中使用了自定义模型，也一并分享到平台，供大家参考。

效果如下：

2. 创建自定义模型

2.1 创建item成员类对象

        创建my_tree_item.py

class MyTreeItem:
    def __init__(self, data, parent=None):
        self.parentItem = parent  # 父节点
        self.childItems = []  # 子节点
        self.itemData = data  # 子节点数据
        self._row = -1  # 此节点位于父节点的位置

    def appendChild(self, child):
        """ 在当前节点插入子节点 """
        child.setRow(len(self.childItems))
        self.childItems.append(child)

    def child(self, row: int):
        """ 获取第row个子节点 """
        return self.childItems[row]

    def parent(self):
        """ 获取父节点 """
        return self.parentItem

    def childCount(self):
        """ 获取子节点数量 """
        return len(self.childItems)

    def columnCount(self):
        """ 统计节点列数 """
        return len(self.itemData)

    def data(self, column):
        """ 获取节点第column列的数据 """
        return self.itemData[column]

    def setRow(self, row: int):
        """ 设置该节点是其父节点的第几个节点 """
        self._row = row

    def row(self):
        """ 获取该节点是父节点的第几个子节点 """
        return self._row

    def setData(self, data, column):
        self.itemData[column] = data

2.2 创建树模型

        创建 my_tree_model.py

from PyQt5 import QtGui
from PyQt5.QtCore import QAbstractItemModel, QVariant, QModelIndex, Qt

from .my_tree_item import MyTreeItem


class MyTreeModel(QAbstractItemModel):
    def __init__(self, parent=None):
        super(MyTreeModel, self).__init__(parent)
        self.rootItem = None  # 最顶层根节点
        self.updateData()

    def data(self, index: QModelIndex, role: int = ...):
        """
        获取索引元素的数据
        :param index:
        :param role:
        :return:
        """
        if not index.isValid():
            return QVariant()

        # 在首列添加icon
        if role == Qt.DecorationRole and index.column() == 0:
            icon = QtGui.QIcon("./image.png")
            return icon
        # 树成员显示的数据
        item = index.internalPointer()
        if role == Qt.DisplayRole and index.column() == 0:
            return item.data(0)

        return QVariant()

    def setData(self, index: QModelIndex, value: typing.Any, role: int = ...) -> bool:
        if role == Qt.QEditRole:
            value_int = value

            self.rootItem[index.row()] = value_int

            self.dataChanged.emit(index, index)

            return True

    def flags(self, index: QModelIndex):
        """
        编辑
        :param index:
        :return:
        """
        if not index.isValid():
            return 0
        return Qt.ItemIsEnabled | Qt.ItemIsSelectable

    def headerData(self, section: int, orientation: Qt.Orientation, role: int = Qt.DisplayRole) -> QVariant:
        """
        标题的显示
        :param section:
        :param orientation:
        :param role:
        :return:
        """
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self.rootItem.data(section)
        return QVariant()

    def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> QModelIndex:
        """
        构造父节点下子节点索引
        :param row:
        :param column:
        :param parent:
        :return:
        """
        if not self.hasIndex(row, column, parent):
            return QModelIndex()

        if not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()

        childItem = parentItem.child(row)

        if childItem:
            return self.createIndex(row, column, childItem)
        else:
            return QModelIndex()

    def parent(self, index: QModelIndex) -> QModelIndex:
        """
        通过子节点获取父节点索引
        :param index:
        :return:
        """
        if not index.isValid():
            return QModelIndex()

        childItem = index.internalPointer()
        parentItem = childItem.parent()

        if parentItem == self.rootItem:
            return QModelIndex()

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

    def rowCount(self, parent: QModelIndex = QModelIndex()) -> int:
        """
        获取父节点下子节点行数
        :param parent:
        :return:
        """
        if not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()

        return parentItem.childCount()

    def columnCount(self, parent: QModelIndex = QModelIndex()) -> int:
        """
        获取父节点下子节点列数
        :param parent:
        :return:
        """
        return self.rootItem.columnCount()

    def setupModelData(self, parent):
        for i in range(5):
            primaryData = [f"节点{i}", "", ""]
            primaryItem = MyTreeItem(primaryData, parent)

            for j in range(3):
                childData = [f"子节点{j}", "", ""]
                childItem = MyTreeItem(childData, primaryItem)
                primaryItem.appendChild(childItem)

            parent.appendChild(primaryItem)

    def updateData(self):
        if self.rootItem:
            del self.rootItem
            self.rootItem = None

        rootData = ["根节点", ""]

        self.rootItem = MyTreeItem(rootData)
        self.setupModelData(self.rootItem)

2.3 结合模型添加按钮

from PyQt5.QtWidgets import QTreeView, QPushButton
from PyQt5.QtCore import QModelIndex

from .my_tree_model import MyTreeModel


class MyTreeView(QTreeView):
    def __init__(self, parent=None):
        super(MyTreeView, self).__init__(parent)
        """ 模型 """
        self.source_model = MyTreeModel()

        self.init()

    def init(self):

        self.setModel(self.source_model)

        self.setColumnWidth(0, 224)
        self.setColumnWidth(1, 40)
        self.initMenu()

    def initMenu(self) -> None:
        # 统计主目录数
        count_primary = self.source_model.rootItem.childCount()
        # 遍历主目录
        for i in range(count_primary):
            # 获取主目录索引
            index = self.source_model.index(i, 1, QModelIndex())
            item = self.source_model.rootItem.child(i)
            self.addMenuButton(item, index)

    def addMenuButton(self, treeItem, index):
        """

        :param treeItem:
        :param index:
        :param row:
        :return:
        """
        # 添加菜单按钮
        button = QPushButton("more")
        self.setIndexWidget(index, button)

        child_count = treeItem.childCount()
        if child_count != 0:
            for i in range(child_count):
                child_item = treeItem.child(i)
                child_index = self.source_model.index(i, 1, index)
                self.addMenuButton(child_item, child_index)

2.4 验证 

import sys
from PyQt5 import QtWidgets
from .my_tree_view import MyTreeView

if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    MainWindow = MyTreeView()
    MainWindow.show()
    sys.exit(app.exec_())