QTreeView: setting icon for item on dropEvent()

问题

I have a QTreeView class and want to set an icon on the dropEvent() event. See the following example:

See the icon for the parent "Users" under the first root "Master Data". I want to drag a user (e.g. "user_a") under the second root "Security Model" under "client_c"/"stakeholder_d" and the user_a should get the same icon as the group (already the case for existing "user_b").

#!/usr/bin/env python3
# coding = utf-8
from PyQt5 import QtWidgets, QtCore, QtGui

TXT_CLIENT = "Clients"
TXT_STAKEHLD = "Stakeholders"
TXT_USER = "Users"

TXT_SYSTEM = "Master Data"
TXT_SECURITY = "Security Model"

CLS_LVL_ROOT = 0
CLS_LVL_CLIENT = 1
CLS_LVL_STAKEHLD = 2
CLS_LVL_USER = 3

ICON_LVL_CLIENT = "img/icons8-bank-16.png"
ICON_LVL_STAKEHLD = "img/icons8-initiate-money-transfer-24.png"
ICON_LVL_USER = "img/icons8-checked-user-male-32.png"

DATA = [
    (TXT_SYSTEM, [
    (TXT_USER, [
        ("user_a", []),
        ("user_b", [])
        ]),
    (TXT_CLIENT, [
        ("client_a", []),    
        ("client_b", []),    
        ("client_c", []),
        ("client_d", [])
        ]),
    (TXT_STAKEHLD, [
        ("stakeholder_a", []),    
        ("stakeholder_b", []),    
        ("stakeholder_c", []),            
        ("stakeholder_d", [])        
        ])
    ]),
    (TXT_SECURITY, [
        ("client_a", [
            ("stakeholder_b",[
                ("user_a",[])
            ])
        ]),
        ("client_c", [
            ("stakeholder_d",[
                ("user_b",[])
            ])
        ])
    ])
    ]

def create_tree_data(tree):
    model = QtGui.QStandardItemModel()
    addItems(tree, model, DATA)
    tree.setModel(model)

def addItems(tree, parent, elements, level=0, root=0):
    level += 1

    for text, children in elements:                    
        if text == TXT_SYSTEM:
            root = 1

        elif text == TXT_SECURITY:
            root = 2

        item = QtGui.QStandardItem(text)            
        icon = QtGui.QIcon(TXT_USER)
        item.setIcon(icon)

        parent.appendRow(item)

        if root==1:
            if children:
                if text == TXT_CLIENT:                    
                    icon = QtGui.QIcon(ICON_LVL_CLIENT)
                    item.setIcon(icon)
                elif text == TXT_STAKEHLD:                                        
                    icon = QtGui.QIcon(ICON_LVL_STAKEHLD)
                    item.setIcon(icon)
                elif text == TXT_USER:                                                            
                    icon = QtGui.QIcon(ICON_LVL_USER)
                    item.setIcon(icon)
        elif root == 2:
            if level == 2:
                icon = QtGui.QIcon(ICON_LVL_CLIENT)
                item.setIcon(icon)
            if level == 3:
                icon = QtGui.QIcon(ICON_LVL_STAKEHLD)
                item.setIcon(icon)
            elif level == 4:
                icon = QtGui.QIcon(ICON_LVL_USER)
                item.setIcon(icon)

        addItems(tree, item, children, level, root)

def get_tree_selection_level(index):
    level = 0
    while index.parent().isValid():
        index = index.parent()
        level += 1

    return level


class TreeView(QtWidgets.QTreeView):
    def __init__(self, parent=None):
        super().__init__(parent)

        self.initUI()

    def initUI(self):
        self.setHeaderHidden(True)
        self.setColumnHidden(1, True)
        self.setSelectionMode(self.SingleSelection)
        self.setDragDropMode(QtWidgets.QAbstractItemView.DragDrop)  # InternalMove)        

    def dropEvent(self, event):
        tree = event.source()        

        if self.viewport().rect().contains(event.pos()):
            fake_model = QtGui.QStandardItemModel()
            fake_model.dropMimeData(
                event.mimeData(), event.dropAction(), 0, 0, QtCore.QModelIndex()
            )            
            for r in range(fake_model.rowCount()):
                for c in range(fake_model.columnCount()):
                    ix = fake_model.index(r, c)
                    print("item: ", ix.data())

                    item = QtGui.QStandardItem(ix.data())
                    icon = QtGui.QIcon(TXT_USER)
                    item.setIcon(icon)

            sParent: str = ""
            par_ix = tree.selectedIndexes()[0].parent()
            if par_ix.isValid():
                sParent = par_ix.data()
                print("par. item: ", sParent)

            to_index = self.indexAt(event.pos())
            if to_index.isValid():
                print("to:", to_index.data())

            if (sParent == TXT_CLIENT and get_tree_selection_level(to_index) == CLS_LVL_ROOT) or (sParent == TXT_STAKEHLD and get_tree_selection_level(to_index) == CLS_LVL_CLIENT) or (sParent == TXT_USER and get_tree_selection_level(to_index) == CLS_LVL_STAKEHLD):
                # to-do:
                # 1 - check if the item is already there; if yes: omit
                pass

                # 2 - set the proper icon
                pass

                super().dropEvent(event)
                self.setExpanded(to_index, True)

class MainWindow(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()

        self.initUI()        

    def initUI(self):
        centralwidget = QtWidgets.QWidget()
        self.setCentralWidget(centralwidget)

        hBox = QtWidgets.QHBoxLayout(centralwidget)        

        self.treeView = TreeView(centralwidget)       

        hBox.addWidget(self.treeView)


if __name__ == "__main__":
    app = QtWidgets.QApplication([])
    window = MainWindow()
    create_tree_data(window.treeView)
    window.treeView.expand(window.treeView.model().index(0, 0))  # expand the System-Branch
    window.setGeometry(400, 400, 500, 400)
    window.show()

    app.exec_()

回答1:

Considering my previous answer you can override the data() method to return the icon with respect to the node level:

class CustomModel(QtGui.QStandardItemModel):
    # ...

    def data(self, index, role=QtCore.Qt.DisplayRole):
        if role == QtCore.Qt.DecorationRole:
            if index == self.item(1).index():
                return
            root = index
            level = -1
            while root.parent().isValid():
                root = root.parent()
                level += 1
            if root == self.item(1).index():
                icon_path = (ICON_LVL_CLIENT, ICON_LVL_STAKEHLD, ICON_LVL_USER)[level]
                icon = QtGui.QIcon(os.path.join(CURRENT_DIR, icon_path))
                return icon
        return super().data(index, role)

来源：https://stackoverflow.com/questions/59530061/qtreeview-setting-icon-for-item-on-dropevent