class StatTreeModel(QAbstractItemModel): RoleId = Qt.UserRole + 1 RoleTier = RoleId + 1 RoleProgress = RoleTier + 1 TIER_1, \ TIER_2, \ TIER_3 = range(1, 4) ColumnProgress = 3 def __init__(self, parent=None, headers=None): super().__init__(parent) self._headers = headers or list() self._rootNode = TreeNode(None, None) def init(self, rootNode): self.beginResetModel() self._rootNode = rootNode self.endResetModel() def clear(self): self._rootNode.clear() def headerData(self, section, orientation, role=None): if orientation == Qt.Horizontal and role == Qt.DisplayRole and section < len( self._headers): return QVariant(self._headers[section]) return QVariant() def rowCount(self, parent=QModelIndex(), *args, **kwargs): if parent.column() > 0: return 0 parent_node = parent.internalPointer() if parent.isValid( ) else self._rootNode return parent_node.child_count() def columnCount(self, parent=QModelIndex(), *args, **kwargs): return len(self._headers) def index(self, row, col, parent): if not self.hasIndex(row, col, parent): return QModelIndex() parent_node = parent.internalPointer() if parent.isValid( ) else self._rootNode child_node = parent_node.child_at(row) return self.createIndex(row, col, child_node) if child_node else QModelIndex() def parent(self, index: QModelIndex): if not index.isValid(): return QModelIndex() child_node = index.internalPointer() if not child_node: return QModelIndex() parent_node = child_node.parent() if parent_node == self._rootNode or parent_node is None: return QModelIndex() return self.createIndex(parent_node.row(), index.column(), parent_node) def data(self, index: QModelIndex, role=None): if not index.isValid(): return QVariant() row = index.row() col = index.column() node = index.internalPointer() if role == Qt.DisplayRole or role == Qt.ToolTipRole: if col == 0: return QVariant(row + 1) return QVariant(node[col]) elif role == self.RoleId: return node['id'] elif role == self.RoleTier: return node['tier'] elif role == self.RoleProgress: if col == self.ColumnProgress: if node['tier'] == self.TIER_1: return node['progress'] elif role == Qt.BackgroundRole: retcol = QColor(Qt.white) if col == self.ColumnProgress: if node['tier'] == self.TIER_2: if node['received']: retcol = QColor(Qt.green).lighter(150) else: retcol = QColor(Qt.red).lighter(150) return QBrush(retcol) return QVariant() def flags(self, index: QModelIndex): f = super().flags(index) return f