Tracking checked items in QTreeWidget when loading in new set of data

Question

I have a QTreeWidget in my gui in which the contents will be cleared whenever it loads in a different set of data and I am trying to tracked what has been checked as User loads in different data set.

Initially, I thought of tracking it using derive_tree_items method that I created in which it contains the QTreeWidgetItem object, however as soon as I tried to load in a new set of data, the objects that I stored will be lost as they are deleted (expected)..

Currently at a lost what is a better way to 'track' these checkable items? (I may also need to populate them into QMenu + QAction, hence the trackable checking but that will be for next time)

In my code, you can replicate by:

• Click on button 'Data-01'
• Check any objects, eg. I checked 'c102' and 'a102'
• Click on button 'Data-02'
• Click on button 'Data-01' again
• Expecting to see 'c102', 'a102' is checked..

IsNewItemRole = QtCore.Qt.UserRole + 1000

class CustomTreeWidgetItem(QtGui.QTreeWidgetItem):
    """Initialization class for QTreeWidgetItem creation.

    Args:
        widget (QtGui.QTreeWidget): To append items into.
        text (str): Input name for QTreeWidgetItem.
        is_tristate (bool): Should it be a tri-state checkbox. False by default.
    """
    def __init__(self, parent=None, text=None, is_tristate=False, is_new_item=False):
        super(CustomTreeWidgetItem, self).__init__(parent)

        self.setText(0, text)
        # flags = QtCore.Qt.ItemIsEditable | QtCore.Qt.ItemIsUserCheckable

        if is_tristate:
            # flags |= QtCore.Qt.ItemIsTristate

            # Solely for the Parent item
            self.setFlags(
                self.flags()
                | QtCore.Qt.ItemIsTristate
                | QtCore.Qt.ItemIsEditable
                | QtCore.Qt.ItemIsUserCheckable
            )
        else:
            self.setFlags(
                self.flags()
                | QtCore.Qt.ItemIsEditable
                | QtCore.Qt.ItemIsUserCheckable
            )
            self.setCheckState(0, QtCore.Qt.Unchecked)

        self.setData(0, IsNewItemRole, is_new_item)

    def setData(self, column, role, value):
        """Override QTreeWidgetItem setData function.

        QTreeWidget does not have a signal that defines when an item has been
        checked/ unchecked. And so, this method will emits the signal as a
        means to handle this.

        Args:
            column (int): Column value of item.
            role (int): Value of Qt.ItemDataRole. It will be Qt.DisplayRole or
                Qt.CheckStateRole
            value (int or unicode): 
        """
        state = self.checkState(column)
        QtGui.QTreeWidgetItem.setData(self, column, role, value)
        if (role == QtCore.Qt.CheckStateRole and
                state != self.checkState(column)):
            tree_widget = self.treeWidget()
            if isinstance(tree_widget, CustomTreeWidget):
                tree_widget.itemToggled.emit(self, column)


class CustomTreeWidget(QtGui.QTreeWidget):
    """Initialization class for QTreeWidget creation.

    Args:
        widget ():
    """
    # itemToggled = QtCore.pyqtSignal(QtGui.QTreeWidgetItem, bool)
    itemToggled = QtCore.Signal(QtGui.QTreeWidgetItem, bool)

    contentUpdates = QtCore.Signal()

    def __init__(self, widget=None):
        super(CustomTreeWidget, self).__init__(widget)

        self.rename_counter = False

        # self.itemToggled.connect(self.handleItemToggled)
        self.currentItemChanged.connect(self.selection_item_changed)
        self.itemChanged.connect(self.tree_item_changed)
        self.itemDoubleClicked.connect(self.tree_item_double_clicked)

    def selection_item_changed(self, current, previous):
        """Overrides widget's default signal.

        Emiited when current item selection is changed. This will also toggles
        the state of `self.add_child_btn`.
        If a child item is selected, the "Add Child" button will be disabled.

        Args:
            current (CustomTreeWidgetItem): Currently selected item.
            previous (CustomTreeWidgetItem or None): Previous selected item.
        """
        state = True
        if not current or current.parent():
            state = False

    def tree_item_changed(self, item, column):
        """Overrides widget's default signal.

        Emitted when the contents of the selected item in the column changes.

        Args:
            item (CustomTreeWidgetItem): Selected item.
            column (int): Column value of the selected item.
        """
        if self.rename_counter and self.prev_name != item.text(column):
            self.rename_counter = False
            item.setData(0, IsNewItemRole, True)

            self.contentUpdates.emit()

        elif item.checkState(column) == QtCore.Qt.Checked:
            print('Item Checked')

        elif item.checkState(column) == QtCore.Qt.Unchecked:
            print('Item Unchecked')

    def tree_item_double_clicked(self, item, column):
        """Overrides widget's default signal.

        Emitted when User performs double clicks inside the widget.

        Args:
            item (CustomTreeWidgetItem): Selected item.
            column (int): Column value of the selected item.
        """
        self.prev_name = item.text(column)
        self.rename_counter = True

    def derive_tree_items(self, mode="all"):
        all_items = OrderedDict()

        root_item = self.invisibleRootItem()
        top_level_count = root_item.childCount()

        for i in range(top_level_count):
            top_level_item = root_item.child(i)
            top_level_item_name = str(top_level_item.text(0))
            child_num = top_level_item.childCount()

            all_items[top_level_item_name] = []

            for n in range(child_num):
                child_item = top_level_item.child(n)
                child_item_name = str(child_item.text(0)) or ""

                all_items[top_level_item_name].append(child_item)

        return all_items


class MainApp(QtGui.QWidget):
    def __init__(self, parent=None):
        super(MainApp, self).__init__(parent)

        self._diff_highlight = False
        self._tree = CustomTreeWidget()
        self._tree.header().hide()

        # QTreeWidget default signals override
        self._tree.contentUpdates.connect(self.update_dictionary)

        tree_layout = QtGui.QVBoxLayout()
        self.btn1 = QtGui.QPushButton("Data-01")
        self.btn2 = QtGui.QPushButton("Data-02")
        tree_layout.addWidget(self._tree)
        tree_layout.addWidget(self.btn1)
        tree_layout.addWidget(self.btn2)

        main_layout = QtGui.QHBoxLayout()
        main_layout.addLayout(tree_layout)
        self.setLayout(main_layout)

        self.setup_connections()

    def setup_connections(self):
        self.btn1.clicked.connect(self.show_data_01)
        self.btn2.clicked.connect(self.show_data_02)

    def update_dictionary(self):
        print '>>> update: ', self._tree.derive_tree_items()

    def show_data_01(self):
        print '>>> Button1 test'

        self._tree.clear()

        test_dict1 = {
            "itemA" :{
                "menuA": ["a101", "a102"],
            },
            "itemBC": {
                "menuC": ["c101", "c102", "c103"],
                "menuB": ["b101"]
            },
        }

        for page_name, page_contents in test_dict1.items():
            # page_item = PageHeaderItem(self._tree, page_name)
            for pk, pv in page_contents.items():
                parent = CustomTreeWidgetItem(self._tree, pk, is_tristate=True)
                for c in pv:
                    child = CustomTreeWidgetItem(parent, c)

        self._tree.expandAll()



    def show_data_02(self):
        print '>>> Button2 test'

        self._tree.clear()

        test_dict2 = {
            "itemD" :{
                "menuD": ["d100"],
            },
        }

        for page_name, page_contents in test_dict2.items():
            # page_item = PageHeaderItem(self._tree, page_name)
            for pk, pv in page_contents.items():
                parent = CustomTreeWidgetItem(self._tree, pk, is_tristate=True)
                for c in pv:
                    child = CustomTreeWidgetItem(parent, c)

        self._tree.expandAll()



if __name__ == "__main__":
    app = QtGui.QApplication(sys.argv)
    w = MainApp()
    w.show()
    sys.exit(app.exec_())

Answer 1

A QTreeWidget (like QListWidget and QTableWidget) has its internal model; it's some sort of a high-level access to a data model, and its actual model is not directly (as in easily) accessible, nor it should. They are "simplified" model-view interfaces, intended for general use that don't require advanced editing, but - most importantly - they only support their own, single and unique model. There's no easy way to change it except from their Q[viewType]WidgetItem interfaces, unless you completely reset the model, meaning that you'll need to "store" the data somewhere else if you want to use multiple models in the same view, making the whole thing much more complex than it needs to be and much prone to errors and issues, which is exactly what happens in your case.

On the opposite side, those QWidgetItemViews offer some features missing in standard models and views, and one of them is the "auto-check" of items in QTreeWidgets.
While that feature is very useful, it could be a serius PITA when you need to show different data models on the same view; this means that, to avoid the rhetorical reinvention of the wheel, it's better to stick with the QTreeView/QStandardItemModel pair and just implement the tristate mechanism instead of using convoluted methods that might clash with the internal implementation of a QTreeWidget.

Separate QStandardItemModel subclass instancies, with parent/child Tristate support

The most important aspect here is that you will be using a single data model class instance for each data set (instead of multiple dict + view's model pairs), making it much easier to switch between them with a simple flick of setModel().
The drawback is the aforementioned lack of parent/children state support, which has to be implemented; once that logic is solved, you'll get multiple persistent, unique and consistent models, no matter how many of them you actually need.

Besides the actual model contents initialization, you are only required to subclass two methods of QStandardItemModel:

• setData(index, value, role) is overridden to apply the check state to the children indexes: if the role is Qt.CheckState and the index has any children, the [un]checked state is applied to them; if the index has a parent, the index emits the dataChanged signal to the model, ensuring that its view[s] requires updates (otherwise the checkbox visible state won't be updated correctly until the view is repainted)^[1];
• data(index, role) overriding is required to "show" the checkstate for the parent(s); it doesn't matter what the model's index data is: if it has any children, its state totally depends on them (all/any/none checked), otherwise it's based on the default model index's checkState;

Once that's solved, you only have to care about setting the newly selected model to the view, and all states will be there as they were before switching to another model, if any.

To keep some consistence with your example, I used your dict-based model data creation logic, but I'd suggest you to use a recursive method to add sub-children.

Since I was already there, I also added a mechanism to store the expanded state of every index, for better view/model consistency; it's not required, but it really helps user experience :-) Keep in mind that that's just there for demonstration purposes: obviously, if you add/remove items without taking care of the internal expandState dict, this won't work properly (or won't work at all!).

import sys
from PyQt5 import QtCore, QtGui, QtWidgets

dataSets = [
    {
        "itemA" :{
            "menuA": ["a101", "a102"],
        },
        "itemBC": {
            "menuC": ["c101", "c102", "c103"],
            "menuB": ["b101"]
        },
    }, 
    {
        "itemD" :{
            "menuD": ["d100"],
        },
    }

]

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

        # unserialize data, as per your original code; you might want to use a
        # recursive function instead, to allow multiple levels of items
        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)

    def setData(self, index, value, role=QtCore.Qt.EditRole):
        if role == QtCore.Qt.CheckStateRole:
            childState = QtCore.Qt.Checked if value else QtCore.Qt.Unchecked
            # set all children states according to this 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)
            # if the item has a parent, emit the dataChanged signal to ensure
            # that the parent state is painted correctly according to what data()
            # will return; note that this will emit the dataChanged signal whatever
            # the "new" parent state is, meaning that it might still be the same
            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 doesn't support auto tristate based on its children 
        # as it does for QTreeWidget's internal model; we have to implement that
        if role == QtCore.Qt.CheckStateRole and self.flags(index) & QtCore.Qt.ItemIsTristate:
            childStates = []
            # collect all child check states
            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)
                    # if the state of a children is partially checked we can
                    # stop here and return a partially checked state
                    if childState == QtCore.Qt.PartiallyChecked:
                        return QtCore.Qt.PartiallyChecked
                    childStates.append(childState)
            if all(childStates):
                # all children are checked, yay!
                return QtCore.Qt.Checked
            elif any(childStates):
                # only some children are checked...
                return QtCore.Qt.PartiallyChecked
            # no item is checked, so bad :-(
            return QtCore.Qt.Unchecked
        return super(TreeModel, self).data(index, role)

    def checkStateChange(self, topLeft, bottomRight):
        # if you need some control back to your data outside the model, here is
        # the right place to do it; note that *usually* the topLeft and 
        # bottomRight indexes are the same, expecially with QStandardItemModels
        # but that would not be the same in some special cases
        pass


class Window(QtWidgets.QWidget):
    def __init__(self):
        QtWidgets.QWidget.__init__(self)
        layout = QtWidgets.QGridLayout()
        self.setLayout(layout)

        self.treeView = QtWidgets.QTreeView()
        layout.addWidget(self.treeView)

        self.models = []
        self.expandStates = {}

        for i, dataSet in enumerate(dataSets):
            model = TreeModel(dataSet)
            button = QtWidgets.QPushButton('Data-{:02}'.format(i + 1))
            layout.addWidget(button)
            button.clicked.connect(lambda _, model=model: self.setModel(model))

    def getExpandState(self, expDict, model, index=QtCore.QModelIndex()):
        # set the index expanded state, if it's not the root index:
        # the root index is not a valid index!
        if index.isValid():
            expDict[index] = self.treeView.isExpanded(index)
        # if the index (or root index) has children, set their states
        for row in range(model.rowCount(index)):
            for col in range(model.columnCount(index)):
                childIndex = model.index(row, col, index)
                # if the current index has children, set their expand state
                # using this function, which is recursive
                for childRow in range(model.rowCount(childIndex)):
                    self.getExpandState(expDict, model, childIndex)

    def setModel(self, model):
        if self.treeView.model():
            if self.treeView.model() == model:
                # the model is the same, no need to update anything
                return
            # save the expand states of the current model before changing it
            prevModel = self.treeView.model()
            self.expandStates[prevModel] = expDict = {}
            self.getExpandState(expDict, prevModel)
        self.treeView.setModel(model)
        if model in self.expandStates:
            # if the new model has expand states saved, restore them
            for index, expanded in self.expandStates.get(model, {}).items():
                self.treeView.setExpanded(index, expanded)
        else:
            self.treeView.expandAll()

if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    w = Window()
    w.show()
    sys.exit(app.exec_())

---

^[1]: In this example the dataChanged signal is emitted whenever any child item check state changes. This isn't a big issue, but if you really need to avoid unnecessary dataChanged notifications you might need to add a QtCore.QTimer.singleshot delayed dataChanged signal emission only if the parent state has changed. It's not that hard to achieve, but I didn't think it was really necessary for this example.