PyQt: How to keep QTreeView nodes correctly expanded after a sort

Question

I'm writing a simple test program using QTreeModel and QTreeView for a more complex project later on. In this simple program, I have data in groups which may be contracted or expanded, as one would expect in a QTreeView. The data may also be sorted by the various data columns (QTreeView.setSortingEnabled is True). Each tree item is a list of data, so the sort function implemented in the TreeModel class uses the built-in python list sort:

self.layoutAboutToBeChanged.emit()  
self.rootItem.childItems.sort(key=lambda x: x.itemData[col], reverse=order)  
for item in self.rootItem.childItems:  
    item.childItems.sort(key=lambda x: x.itemData[col], reverse=order)  
self.layoutChanged.emit()

The problem is that whenever I change the sorting of the root's child items (the tree is only 2 levels deep, so this is the only level with children) the nodes aren't necessarily expanded as they were before. If I change the sorting back without expanding or collapsing anything, the nodes are expanded as before the sorting change.
Can anyone explain to me what I'm doing wrong? I suspect it's something with not properly reassigning QModelIndex with the sorted nodes, but I'm not sure.

Answer 1

I solved the problem by implementing a QSortFilterProxyModel for the sort.