Re: QTreeView + sort + delete = crash

from PyQt5 import QtCore, QtWidgets, QtGui
from PyQt5 import uic

from itertools import groupby
import sys



# ------------------------------------------------------------------------------
# | CategoryNode                                                               |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | Defines the nodes of the data tree that underlies CategoryTreeModel        |
# ------------------------------------------------------------------------------
class CategoryNode:
    def __init__(self, desc, parent=None):
        self.parent_node = parent
        self._desc = desc
        self.children = []

    # ----------------------------------------------------------------------
    # | CategoryNode accessor functions
    # ----------------------------------------------------------------------

    # Number of columns is always one in this model
    @staticmethod
    def n_columns():
        return 1

    # Return category description
    def desc(self):
        return self._desc

    # Set category description
    def set_desc(self, desc):
        self._desc = desc

    # Return row 
    def row(self):
        if self.parent_node:
            return self.parent_node.children.index(self)
        else:
            return 0

    # Return parent
    def parent(self):
        return self.parent_node

    # Add a child
    def append_child(self, c):
        self.children.append(c)

    # Return child from given row
    def child(self, row):
        return self.children[row]

    # Return number of children
    def n_children(self):
        return len(self.children)

    # Sort children
    def sort_children(self):
        self.children.sort(key=lambda x: x.desc())
        return

    # Delete child from given row
    def del_child(self, row):
        del self.children[row]

    # Printable representation
    def __repr__(self):
        return str(f'->{self._desc}')

    # Recursively dump tree contents
    def dump(self, i=0):
        print(f'''{(' ' * (i * 4))} {self.row()} -> {self}''')
        for c in self.children:
            c.dump(i+1)


# ------------------------------------------------------------------------------
# | CategoryTreeModel                                                          |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | Model used by CategoryDialog tree view                                     |
# ------------------------------------------------------------------------------
class CategoryTreeModel(QtCore.QAbstractItemModel):

    # ------------------------------------------------------------[override]
    # | __init__()
    # ----------------------------------------------------------------------
    def __init__(self, view, parent=None):
        super().__init__(parent)
        self.view = view

        category_data = [
            (4, 'Auto', None), (10, 'Auto', 'Gas'), (9, 'Auto', 'License'),
                (11, 'Auto', 'Service'),
            (23, 'Beer', None),
            (12, 'Clothing', None), (13, 'Clothing', 'Cleaning'),
            (18, 'Dining', None), (19, 'Dining', 'Breakfast'), (21, 'Dining', 'Lunch'),
            (14, 'Healthcare', None), (26, 'Healthcare', 'Dental'),
            (3, 'Household', None),
            (5, 'Pet Expenses', None), (7, 'Pet Expenses', 'Biskies'),
                (17, 'Pet Expenses', 'Chow'), (8, 'Pet Expenses', 'Grooming'),
                (6, 'Pet Expenses', 'Medical'), (15, 'Pet Expenses', 'Miscellaneous'),
            (25, 'Reimbursible', None)
        ]

        # Create root node
        self.root = CategoryNode('Category')

        # Create main category/subcategory node structure
        for g in groupby(category_data, lambda rec: rec[1]):
            sub_list = list(g[1])
            main_node = CategoryNode(g[0], self.root)
            self.root.append_child(main_node)
            for s in sub_list:
                if s[2] is not None:
                    sub_node = CategoryNode(s[2], main_node)
                    main_node.append_child(sub_node)

    # ------------------------------------------------------------[override]
    # | rowCount()
    # ----------------------------------------------------------------------
    def rowCount(self, parent=QtCore.QModelIndex()):
        if parent.isValid():
            return parent.internalPointer().n_children()
        else:
            return self.root.n_children()

    # ------------------------------------------------------------[override]
    # | columnCount()
    # ----------------------------------------------------------------------
    def columnCount(self, parent=QtCore.QModelIndex()):
        if parent.isValid():
            return parent.internalPointer().n_columns()
        else:
            return self.root.n_columns()

    # ------------------------------------------------------------[override]
    # | data()
    # ----------------------------------------------------------------------
    def data(self, index, role=QtCore.Qt.DisplayRole):
        if index.isValid():
            if role == QtCore.Qt.DisplayRole:
                return index.internalPointer().desc()
            elif role == QtCore.Qt.FontRole:
                font = QtGui.QFont('Candara')
                font.setPointSize(11)
                return font
            elif role == QtCore.Qt.ForegroundRole:
                color = '#1e62ce'
                return QtGui.QBrush(QtGui.QColor(color))

        return None

    # ------------------------------------------------------------[override]
    # | flags()
    # ----------------------------------------------------------------------
    def flags(self, index):
        if index.isValid():
            return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable
        else:
            return QtCore.Qt.NoItemFlags

    # ------------------------------------------------------------[override]
    # | index()
    # ----------------------------------------------------------------------
    def index(self, row, column, parent=QtCore.QModelIndex()):

        # According to legend, index() may sometimes be called with invalid parameters
        # (that is, with parameters that don't actually have an index).  So this is
        # necessary to keep it from crashing.
        # https://stackoverflow.com/questions/26680168
        #     /pyqt-treeview-index-error-removing-last-row
        if not self.hasIndex(row, column, parent):
            return QtCore.QModelIndex()

        elif parent.isValid():
            return self.createIndex(row, column, parent.internalPointer().child(row))
        else:
            return self.createIndex(row, column, self.root.child(row))

    # ------------------------------------------------------------[override]
    # | parent()
    # ----------------------------------------------------------------------
    def parent(self, index):
        if not index.isValid():
            return QtCore.QModelIndex()

#       print(f'parent(), before index.internalPointer() call - row={index.row()}')
        child_node = index.internalPointer()
#       print(f'parent(), after index.internalPointer() call  - child_node={child_node}')
        parent_node = child_node.parent()
        if parent_node == self.root:
            r = QtCore.QModelIndex()
        else:
            r = self.createIndex(parent_node.row(), 0, parent_node)
        return r


    # ------------------------------------------------------------[override]
    # | setData()
    # ----------------------------------------------------------------------
    # noinspection PyUnresolvedReferences
    def setData(self, index, value, role=QtCore.Qt.EditRole):
        if role == QtCore.Qt.EditRole:

            # Set value in data tree node
            node = index.internalPointer()
            node.set_desc(value)

            # Re-sort at the level of the modified node (i.e., the children of the
            # modified node's parent)
            # Need to emit layoutAboutToBeChanged and layoutChanged signals to get
            # QTreeView to update properly
            self.layoutAboutToBeChanged.emit([], QtCore.QAbstractItemModel.VerticalSortHint)
            node.parent().sort_children()
            self.layoutChanged.emit([], QtCore.QAbstractItemModel.VerticalSortHint)

            # Make modified item current
            parent = self.parent(index)
            self.view.setCurrentIndex(self.index(node.row(), 0, parent))
            self.view.setFocus(QtCore.Qt.OtherFocusReason)


    # ------------------------------------------------------------[override]
    # | insertRows()
    # ----------------------------------------------------------------------
    def insertRows(self, row, rows, parent=QtCore.QModelIndex()):
        assert rows == 1

        # Determine parent
        if parent.internalPointer() is None:
            parent_node = self.root
        else:
            parent_node = parent.internalPointer()

        # Create new data tree node and add it to parent
        new = CategoryNode(None, parent_node)
        self.beginInsertRows(parent, row, row + rows - 1)
        parent_node.append_child(new)
        self.endInsertRows()

        return True

    # ------------------------------------------------------------[override]
    # | removeRows()
    # ----------------------------------------------------------------------
    def removeRows(self, row, rows, parent=QtCore.QModelIndex()):
        assert rows == 1

        # Determine parent
        if parent.internalPointer() is None:
            parent_node = self.root
        else:
            parent_node = parent.internalPointer()

        # Delete it from the parent
        self.beginRemoveRows(parent, row, row + rows - 1)
        parent_node.del_child(row)
        self.endRemoveRows()

        return True


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


# ------------------------------------------------------------------------------
# | Dlg                                                                        |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | Dialog widget (defined in dlg.ui)                                          |
# ------------------------------------------------------------------------------
dlg_base, dlg_form = uic.loadUiType('dlg.ui')
class Dlg(dlg_base, dlg_form):

    def __init__(self, parent=None):
        super(dlg_base, self).__init__(parent)
        self.setupUi(self)
        self.setWindowFlags(QtCore.Qt.Dialog|QtCore.Qt.FramelessWindowHint)
        self.treeView.setHeaderHidden(True)
        self.treeView.setAnimated(True)

        # Set model
        self.model = CategoryTreeModel(self.treeView)
        self.treeView.setModel(self.model)

        # Set Delegate
        self.delegate = CategoryDelegate()
        self.treeView.setItemDelegate(self.delegate)

        # Connect buttons
        self.addMainButton.clicked.connect(self.add_main)
        self.addSubButton.clicked.connect(self.add_sub)
        self.deleteButton.clicked.connect(self.delete)
        self.doneButton.clicked.connect(quit)

    def keyPressEvent(self, event):
        key = event.key()
        mod = int(event.modifiers())
        if key == QtCore.Qt.Key_Q and mod == QtCore.Qt.CTRL:
            exit()

    def add_main(self):
        print('Add Main')
        view = self.treeView
        model = view.model()
        model.insertRows(model.rowCount(), 1)
        index = model.index(model.rowCount() - 1, 0)
        view.setCurrentIndex(index)
        view.scrollTo(index, QtWidgets.QAbstractItemView.EnsureVisible)
        view.edit(index)

    def add_sub(self):
        print('Add Sub')
        view = self.treeView
        model = view.model()
        current = view.currentIndex()
        if model.parent(current).internalPointer():
            add_to = model.parent(current)
        else:
            add_to = current
        model.insertRows(model.rowCount(add_to), 1, add_to)
        index = model.index(model.rowCount(add_to) - 1, 0, add_to)
        view.setCurrentIndex(index)
        view.scrollTo(index, QtWidgets.QAbstractItemView.EnsureVisible)
        view.edit(index)

    def delete(self):
        print('Delete')
        view = self.treeView
        model = view.model()
        current = view.currentIndex()
        model.removeRows(current.row(), 1, model.parent(current))


# ------------------------------------------------------------------------------
# | CategoryDelegate                                                           |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | Delegate for editing                                                       |
# ------------------------------------------------------------------------------
class CategoryDelegate(QtWidgets.QStyledItemDelegate):

    def __init__(self, parent=None):
        super().__init__(parent)

    def createEditor(self, parent, option, index):
        self.editor = QtWidgets.QLineEdit(parent)
        return self.editor

    def setEditorData(self, editor, index):
        pass

    def updateEditorGeometry(self, editor, option, index):
        editor.setGeometry(option.rect)

    def setModelData(self, editor, model, index):
        model.setData(index, editor.text())


# ------------------------------------------------------------------------------
# | pytest tester function                                                     |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# ------------------------------------------------------------------------------
def test_model(qtmodeltester):
    dlg = Dlg()
    qtmodeltester.check(dlg.model)


# ------------------------------------------------------------------------------
# | main()                                                                     |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# ------------------------------------------------------------------------------
if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    app.setStyle('Fusion')
    dlg = Dlg()
    dlg.show()
    sys.exit(app.exec_())

<?xml version="1.0" encoding="UTF-8"?>
<ui version="4.0">
 <class>CategoryDialog</class>
 <widget class="QDialog" name="CategoryDialog">
  <property name="geometry">
   <rect>
    <x>0</x>
    <y>0</y>
    <width>381</width>
    <height>423</height>
   </rect>
  </property>
  <property name="font">
   <font>
    <family>Corbel</family>
    <pointsize>10</pointsize>
   </font>
  </property>
  <property name="windowTitle">
   <string>Dialog</string>
  </property>
  <property name="modal">
   <bool>true</bool>
  </property>
  <widget class="QWidget" name="layoutWidget">
   <property name="geometry">
    <rect>
     <x>10</x>
     <y>10</y>
     <width>361</width>
     <height>401</height>
    </rect>
   </property>
   <layout class="QVBoxLayout" name="verticalLayout">
    <item>
     <widget class="TreeView" name="treeView">
      <property name="font">
       <font>
        <family>Corbel</family>
        <pointsize>10</pointsize>
       </font>
      </property>
      <property name="cursor" stdset="0">
       <cursorShape>PointingHandCursor</cursorShape>
      </property>
     </widget>
    </item>
    <item>
     <spacer name="horizontalSpacer">
      <property name="orientation">
       <enum>Qt::Horizontal</enum>
      </property>
      <property name="sizeHint" stdset="0">
       <size>
        <width>40</width>
        <height>20</height>
       </size>
      </property>
     </spacer>
    </item>
    <item>
     <layout class="QHBoxLayout" name="horizontalLayout">
      <item>
       <widget class="QToolButton" name="addMainButton">
        <property name="font">
         <font>
          <pointsize>9</pointsize>
         </font>
        </property>
        <property name="text">
         <string>Main</string>
        </property>
        <property name="iconSize">
         <size>
          <width>40</width>
          <height>40</height>
         </size>
        </property>
        <property name="toolButtonStyle">
         <enum>Qt::ToolButtonTextUnderIcon</enum>
        </property>
       </widget>
      </item>
      <item>
       <widget class="QToolButton" name="addSubButton">
        <property name="font">
         <font>
          <pointsize>9</pointsize>
         </font>
        </property>
        <property name="text">
         <string>Sub</string>
        </property>
        <property name="iconSize">
         <size>
          <width>40</width>
          <height>40</height>
         </size>
        </property>
        <property name="toolButtonStyle">
         <enum>Qt::ToolButtonTextUnderIcon</enum>
        </property>
       </widget>
      </item>
      <item>
       <widget class="QToolButton" name="deleteButton">
        <property name="font">
         <font>
          <pointsize>9</pointsize>
         </font>
        </property>
        <property name="text">
         <string>Delete</string>
        </property>
        <property name="iconSize">
         <size>
          <width>40</width>
          <height>40</height>
         </size>
        </property>
        <property name="toolButtonStyle">
         <enum>Qt::ToolButtonTextUnderIcon</enum>
        </property>
       </widget>
      </item>
      <item>
       <spacer name="horizontalSpacer_2">
        <property name="orientation">
         <enum>Qt::Horizontal</enum>
        </property>
        <property name="sizeHint" stdset="0">
         <size>
          <width>40</width>
          <height>20</height>
         </size>
        </property>
       </spacer>
      </item>
      <item>
       <widget class="QToolButton" name="doneButton">
        <property name="font">
         <font>
          <pointsize>9</pointsize>
         </font>
        </property>
        <property name="text">
         <string>Done</string>
        </property>
        <property name="iconSize">
         <size>
          <width>40</width>
          <height>40</height>
         </size>
        </property>
        <property name="toolButtonStyle">
         <enum>Qt::ToolButtonTextUnderIcon</enum>
        </property>
       </widget>
      </item>
     </layout>
    </item>
   </layout>
  </widget>
 </widget>
 <customwidgets>
  <customwidget>
   <class>TreeView</class>
   <extends>QTreeView</extends>
   <header>cat.h</header>
  </customwidget>
 </customwidgets>
 <resources/>
 <connections/>
</ui>