Сводная таблица из QTreeView с чекбоксами

Есть код на PyQt5, представляющий собой аналог древовидную таблицу (QTreeView), в которой можно чекбоксами выбирать чекбоксы. Мне нужно чтобы эта таблица работала следующим образом:

  1. если все чекбоксы выключены (не отмечены галогкой) и нажать на родительский класс, то включается родительский и все дочерние - это уже работает, сделано;
  2. если родительский чекбокс все дочерние чекбоксы выключены, и нажать на дочерний, то включается этот дочерний чекбокс и родительский. Остальные же дочерние чекбоксы не включаются;
  3. если все чекбоксы включены (родительские и все дочерние), и включать\выключать один дочерний, то только этот дочерний чекбокс меняется, а остальные - нет.

У меня получилось сделать пункты 1 и 3, а пункт 2 вообще никак не получается. Получалось сделать только так, что если включить дочерний чекбокс (кликнуть по нему), то включаются все дочерние чекбоксы и родительский - но так работать не должно.
Помогите, пожалуйста.

import sys
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import QWidget

import pandas as pd



class Pivot_table(QTreeView):
    def __init__(
            self,
            table_name: str,
            table_width: int,
            columns_width: list,
            data_frame: pd.DataFrame,
            index_cols: list,
            columns_cols: list,
            values_cols: list,
            included_checkBox: bool):
        super().__init__()
        data_frame = data_frame[index_cols]
        pivot_df = data_frame.pivot_table(
            index=index_cols,
            aggfunc='sum',
            fill_value=0)

        header_names = [table_name]
        if columns_cols is not None:
            header_names = header_names + columns_cols
        if values_cols is not None:
            header_names = header_names + values_cols

        self.model_Pivot_table = QStandardItemModel()
        self.model_Pivot_table.setHorizontalHeaderLabels(
            header_names)
        self.model_Pivot_table.itemChanged.connect(
            self.status_checkboxses)

        # Настройка QTreeView
        self.setModel(self.model_Pivot_table)
        self.setFixedWidth(table_width)
        self.setEditTriggers(
            QAbstractItemView.EditTrigger.NoEditTriggers)

        self.populate_model(
            index_cols, pivot_df, included_checkBox)

        if columns_width is not None:  # Установка ширины столбцов
            for column in columns_width:
                self.setColumnWidth(column[0], column[1])

    def populate_model(self, index_cols, pivot_df, included_checkBox):
        parent_items = {}
        for row_index, (index_value, row) in enumerate(pivot_df.iterrows()):
            index_values = index_value \
                if isinstance(index_value, tuple) \
                else (index_value,)
            parent_item = None
            for idx, value in enumerate(index_values):
                if idx == 0:  # Родительский уровень
                    if value not in parent_items:
                        parent_item = QStandardItem(str(value))
                        
                        if included_checkBox:
                            parent_item.setCheckable(True)
                            parent_item.setCheckState(Qt.CheckState.Checked)
                            
                        self.model_Pivot_table.appendRow(parent_item)
                        parent_items[value] = parent_item
                    else:
                        parent_item = parent_items[value]
                
                else:
                    child_found = False
                    for i in range(parent_item.rowCount()):
                        if parent_item.child(i, 0).text() == str(value):
                            child_item = parent_item.child(i, 0)
                            child_found = True
                            break
                    if not child_found:
                        child_item = QStandardItem(str(value))
                        if included_checkBox:
                            child_item.setCheckable(True)
                            child_item.setCheckState(Qt.CheckState.Checked) 
                        parent_item.appendRow(child_item)
                    
                    parent_item = child_item

            # Добавление значения из сводной таблицы
            for col_value in pivot_df.columns:
                value_item = QStandardItem(str(row[col_value]))
                if included_checkBox:
                    value_item.setCheckable(True)
                    value_item.setCheckState(Qt.CheckState.Checked)

                parent_item.appendRow(value_item)

    def status_checkboxses(self, item):
        if not item.parent():
           child_count = item.rowCount()
           for i in range(child_count):
               child_item = item.child(i, 0)
               child_item.setCheckState(item.checkState())


# Основное приложение
if __name__ == '__main__':
    app = QApplication(sys.argv)

    # Данные
    data = {
        0: [1, 2, 3, 3, 4],
        1: [2, 2, 2, 3, 2],
        'ФИО': [
            'Иванов Иван Иванович',
            'Иванов2 Иван2 Иванович2',
            'Иванов3 Иван3 Иванович3',
            'Иванов4 Иван4 Иванович4',
            'Иванов5 Иван5 Иванович5']
    }

    df_residents = pd.DataFrame(data)
    
    widget_apartment = Pivot_table(
        table_name= 'Кв и ФИО',
        table_width= 680,
        columns_width=[[0, 640]],
        data_frame=df_residents,
        index_cols= [1, 'ФИО'],
        columns_cols= None,
        values_cols= None,
        included_checkBox= True)

    # Устанавливаем виджет в главное окно
    central_widget = QWidget()
    layout = QVBoxLayout()
    layout.addWidget(widget_apartment)
    central_widget.setLayout(layout)

    main_window = QMainWindow()
    main_window.setGeometry(100, 100, 800, 600)
    main_window.setCentralWidget(central_widget)
    main_window.show()
    
    sys.exit(app.exec_())

Ответы (1 шт):

Автор решения: S. Nick

Если я правильно понял, вам надо реализовать механизм (с тремя состояниями) Tristate .

Это можно реализовать переопределением двух методов QStandardItemModel:

  • setData(index, value, role) - переопределяется для применения состояния проверки;
  • data(index, role) - требуется для show контрольного состояния для родителя.

main.py

import sys
from PyQt5 import QtCore, QtGui, QtWidgets


class TreeModel(QtGui.QStandardItemModel):
    checkStateChange = QtCore.pyqtSignal(QtCore.QModelIndex, bool)
    
    def __init__(self, dataSet):
        super(TreeModel, self).__init__()

        for page_name, page_contents in dataSet.items():
            for pk, pv in page_contents.items():
                parent = QtGui.QStandardItem(pk)
                parent.setCheckable(True)
                self.appendRow(parent)
                if pv:
                    parent.setTristate(True)
                    for c in pv:
                        child = QtGui.QStandardItem(c)
                        child.setCheckable(True)
                        parent.appendRow(child)

        self.dataChanged.connect(self.checkStateChange)
        self.setHorizontalHeaderLabels(['Кв и ФИО'])

    def setData(self, index, value, role=QtCore.Qt.EditRole):
        if role == QtCore.Qt.CheckStateRole:
            childState = QtCore.Qt.Checked if value \
                                           else QtCore.Qt.Unchecked
            # установить все children состояния в соответствии 
            # с parent item
            for row in range(self.rowCount(index)):
                for col in range(self.columnCount(index)):
                    childIndex = self.index(row, col, index)
                    self.setData(
                        childIndex, childState, 
                        QtCore.Qt.CheckStateRole)
            # если у элемента есть родительский элемент, 
            # подайте сигнал dataChanged, чтобы убедиться, 
            # что родительское состояние отображается правильно 
            # в соответствии с тем, что data() вернет; 
            parent = self.parent(index)
            if parent.isValid():
                self.dataChanged.emit(parent, parent)
        return super(TreeModel, self).setData(index, value, role)

    def data(self, index, role=QtCore.Qt.DisplayRole):
        # QStandardItemModel не поддерживает автоматическое Tristate, 
        # основанное на его дочерних элементах, как это происходит 
        # для внутренней модели QTreeWidget. 
        # Надо реализовать это:
        if role == QtCore.Qt.CheckStateRole and \
            self.flags(index) & QtCore.Qt.ItemIsTristate:
            childStates = []
            # собрать все дочерние состояния проверки
            for row in range(self.rowCount(index)):
                for col in range(self.columnCount(index)):
                    childIndex = self.index(row, col, index)
                    childState = self.data(childIndex, 
                        QtCore.Qt.CheckStateRole)
                    # если состояние child частично проверено, 
                    # мы можем остановиться и 
                    # вернуть частично проверенное состояние
                    if childState == QtCore.Qt.PartiallyChecked:
                        return QtCore.Qt.PartiallyChecked
                    childStates.append(childState)
            if all(childStates):
                # все children проверены
                return QtCore.Qt.Checked
            elif any(childStates):
                # только некоторые children проверены ...
                return QtCore.Qt.PartiallyChecked
            return QtCore.Qt.Unchecked
        return super(TreeModel, self).data(index, role)

    def checkStateChange(self, topLeft, bottomRight):
        pass


# Основное приложение
class MainWindow(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()
        self.central_widget = QtWidgets.QWidget()
        self.setCentralWidget(self.central_widget)
        
        self.treeView = QtWidgets.QTreeView()
        
        layout = QtWidgets.QGridLayout(self.central_widget)
        layout.addWidget(self.treeView)

        model = TreeModel(dataSet)
        self.setModel(model)

    def getExpandState(self, expDict, model, index=QtCore.QModelIndex()):
        # установите expanded state индекса, 
        # если это не корневой индекс: 
        # корневой индекс не является допустимым индексом!
        if index.isValid():
            expDict[index] = self.treeView.isExpanded(index)
        # если у индекса (или корневого индекса) есть дочерние элементы, 
        # установите их состояния
        for row in range(model.rowCount(index)):
            for col in range(model.columnCount(index)):
                childIndex = model.index(row, col, index)
                # если у текущего индекса есть дочерние элементы, 
                # установите их expand state, 
                # используя эту функцию, которая является рекурсивной
                for childRow in range(model.rowCount(childIndex)):
                    self.getExpandState(expDict, model, childIndex)

    def setModel(self, model):
        self.treeView.setModel(model)
        self.treeView.expandAll() 


if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    app.setFont(QtGui.QFont("Times", 12, QtGui.QFont.Bold))

    dataSet = {
        "itemA" :{
            "группа A": ["a101", "a102"],
        },
        "itemBC": {
            "группа B": ["b101"],
            "группа C": ["c101", "c102", "c103"],
        },
        "itemD" :{
            "группа D": ["d100"],
        },
    }    
    
    window = MainWindow()
    window.setWindowTitle("QTreeView Example")
    window.resize(300, 400)    
    window.show()
    sys.exit(app.exec())

введите сюда описание изображения

→ Ссылка