How to implement a QLineEdit in the column header for filtering data inside my QTreeView?

How could I implement a QLineEdit in the column header for filtering data inside my QTreeView?

Expected Result something like this:

So far I've got this:

What I have so far is everything but the QLineEdit for the filter. Full working code example (peewee ORM necessary):

import sys
import re
from peewee import *
from PyQt5 import QtWidgets, QtGui, QtCore 

COUNT_PERS_COLS = 3
col_persID, col_persLAST_NAME, col_persFIRST_NAME = range(COUNT_PERS_COLS)

db = SqliteDatabase(':memory:')

def _human_key(key):
    parts = re.split(r'(\d*\.\d+|\d+)', key)
    return tuple((e.swapcase() if i % 2 == 0 else float(e))
            for i, e in enumerate(parts))

class FilterHeader(QtWidgets.QHeaderView):
    filterActivated = QtCore.pyqtSignal()

    def __init__(self, parent):
        super().__init__(QtCore.Qt.Horizontal, parent)
        self._editors = []
        self._padding = 4
        self.setStretchLastSection(True)
        # self.setResizeMode(QHeaderView.Stretch) obsolete
        self.setSectionResizeMode(QtWidgets.QHeaderView.Stretch)
        self.setDefaultAlignment(
            QtCore.Qt.AlignLeft | QtCore.Qt.AlignVCenter)
        self.setSortIndicatorShown(False)
        self.sectionResized.connect(self.adjustPositions)
        parent.horizontalScrollBar().valueChanged.connect(
            self.adjustPositions)

    def setFilterBoxes(self, count):
        while self._editors:
            editor = self._editors.pop()
            editor.deleteLater()
        for index in range(count):
            editor = QtWidgets.QLineEdit(self.parent())
            editor.setPlaceholderText('Filter')
            editor.returnPressed.connect(self.filterActivated.emit)
            self._editors.append(editor)
        self.adjustPositions()

    def sizeHint(self):
        size = super().sizeHint()
        if self._editors:
            height = self._editors[0].sizeHint().height()
            size.setHeight(size.height() + height + self._padding)
        return size

    def updateGeometries(self):
        if self._editors:
            height = self._editors[0].sizeHint().height()
            self.setViewportMargins(0, 0, 0, height + self._padding)
        else:
            self.setViewportMargins(0, 0, 0, 0)
        super().updateGeometries()
        self.adjustPositions()

    def adjustPositions(self):
        for index, editor in enumerate(self._editors):
            height = editor.sizeHint().height()
            editor.move(
                self.sectionPosition(index) - self.offset() + 2,
                height + (self._padding // 2))
            editor.resize(self.sectionSize(index), height)

    def filterText(self, index):
        if 0 <= index < len(self._editors):
            return self._editors[index].text()
        return ''

    def setFilterText(self, index, text):
        if 0 <= index < len(self._editors):
            self._editors[index].setText(text)

    def clearFilters(self):
        for editor in self._editors:
            editor.clear()


class HumanProxyModel(QtCore.QSortFilterProxyModel):
    def lessThan(self, source_left, source_right):
        data_left = source_left.data()
        data_right = source_right.data()
        if type(data_left) == type(data_right) == str:
            return _human_key(data_left) < _human_key(data_right)
        return super(HumanProxyModel, self).lessThan(source_left, source_right)


class Person(Model):
    persId = CharField()
    lastName = CharField()
    firstName = CharField() 

    class Meta:
        database = db


class winMain(QtWidgets.QMainWindow):
    def __init__(self, parent=None):        
        super().__init__(parent)                
        self.setupUi()
        self.setGeometry(300,200,700,500)
        self.show()

    def createPersonModel(self,parent):        
        model = QtGui.QStandardItemModel(0, COUNT_PERS_COLS, parent)        
        #model.setHeaderData(col_persID, QtCore.Qt.Horizontal, "ID")        
        #model.setHeaderData(col_persLAST_NAME, QtCore.Qt.Horizontal, "Last Name")
        #model.setHeaderData(col_persFIRST_NAME, QtCore.Qt.Horizontal, "First Name")
        #model.setHorizontalHeaderLabels('One Two Three' .split())
        model.setHorizontalHeaderLabels(['ID', 'Last Name', 'First Name'])

        return model

    def addPerson(self, model, id, last_name, first_name):
        model.insertRow(0)        
        model.setData(model.index(0, col_persID), id)
        model.setData(model.index(0, col_persLAST_NAME), last_name)
        model.setData(model.index(0, col_persFIRST_NAME), first_name)

    def handleFilterActivated(self):
        #header = self.tableView.horizontalHeader()  # QTableView()-command
        header = self.treeView.header()
        for index in range(header.count()):
            print((index, header.filterText(index)))                

    def setupUi(self):
        self.treeView = QtWidgets.QTreeView(self)
        self.treeView.setGeometry(0,0,700,500)
        self.treeView.setSortingEnabled(True)
        self.treeView.setAlternatingRowColors(True)        
        self.treeView.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        self.treeView.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
        self.treeView.setAnimated(True)
        self.treeView.setItemsExpandable(True)

        #layout = QtWidgets.QVBoxLayout(self)
        #layout.addWidget(self.treeView)
        header = FilterHeader(self.treeView)
        # self.tableView.setHorizontalHeader(header)  # QTableView()-command
        self.treeView.setHeader(header)        

        model = self.createPersonModel(self)
        self.treeView.setModel(model)

        proxy = HumanProxyModel(self)
        proxy.setSourceModel(model)
        self.treeView.setModel(proxy)

        header.setFilterBoxes(model.columnCount())
        header.filterActivated.connect(self.handleFilterActivated)

        for rec_person in Person.select():
            self.addPerson(model, rec_person.persId, rec_person.lastName, rec_person.firstName)

if __name__ == '__main__':                         
    app = QtWidgets.QApplication(sys.argv)         

    # create a table for our model
    db.create_tables([Person])

    # create some sample data for our model
    Person.create(persId='1001', lastName='Martin', firstName='Robert')
    Person.create(persId='1002', lastName='Smith', firstName='Brad')
    Person.create(persId='1003', lastName='Smith', firstName='Angelina')        

    window = winMain()    
    sys.exit(app.exec_())

I already use the QSortFilterProxyModel class for sorting the result by clicking into the column header.

Solution

You have to override the filterAcceptsRow method by implementing the filtering logic, in the following example it will be filtered using the following rules:

Only the columns that have a different text to the vacuum are considered.
The filtering criterion is that the text placed in the QLineEdit is contained in the corresponding row and column

class HumanProxyModel(QtCore.QSortFilterProxyModel):
    # ...
    @property
    def filters(self):
        if not hasattr(self, "_filters"):
            self._filters = []
        return self._filters

    @filters.setter
    def filters(self, filters):
        self._filters = filters
        self.invalidateFilter()

    def filterAcceptsRow(self, sourceRow, sourceParent):
        for i, text in self.filters:
            if 0 <= i < self.columnCount():
                ix = self.sourceModel().index(sourceRow, i, sourceParent)
                data = ix.data()
                if text not in data:
                    return False
        return True

# ...

class winMain(QtWidgets.QMainWindow):
    # ...
    def handleFilterActivated(self):
        header = self.treeView.header()
        filters = []
        for i in range(header.count()):
            text = header.filterText(i)
            if text:
                filters.append((i, text))
        proxy = self.treeView.model()
        proxy.filters = filters