Setting background on Qtableview using a function

Question

IU have a Qt window with a button and a QTableView which loads some data from a pandas dataframe using the QAbstractTableModel. I would like to set the backgroundcolor of a cel (or a row) when i click on the button.

This is my code so far:

import timeit
import pandas as pd
from PyQt5 import QtWidgets
from PyQt5.QtCore import *
from PyQt5.QtCore import QAbstractTableModel, Qt
from PyQt5.QtGui import *
from PyQt5.uic import loadUi


class PandasModel(QAbstractTableModel):
    def __init__(self, data, parent=None):
        QAbstractItemModel.__init__(self, parent)
        self._data = data
        self.colors = dict()

    def rowCount(self, parent=None):
        return self._data.index.size

    def columnCount(self, parent=None):
        return self._data.columns.size

    def data(self, index, role=Qt.DisplayRole):
        if index.isValid():
            if role == Qt.DisplayRole:
                return str(self._data.iloc[index.row(), index.column()])
            if role == Qt.EditRole:
                return str(self._data.iloc[index.row(), index.column()])
            if role == Qt.BackgroundRole:
                color = self.colors.get((index.row(), index.column()))
                if color is not None:
                    return color
        return None

    def headerData(self, rowcol, orientation, role):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self._data.columns[rowcol]
        if orientation == Qt.Vertical and role == Qt.DisplayRole:
            return self._data.index[rowcol]
        return None

    def flags(self, index):
        flags = super(self.__class__, self).flags(index)
        flags |= Qt.ItemIsEditable
        flags |= Qt.ItemIsSelectable
        flags |= Qt.ItemIsEnabled
        flags |= Qt.ItemIsDragEnabled
        flags |= Qt.ItemIsDropEnabled
        return flags

    def sort(self, Ncol, order):
        """Sort table by given column number.
        """
        try:
            self.layoutAboutToBeChanged.emit()
            self._data = self._data.sort_values(self._data.columns[Ncol], ascending=not order)
            self.layoutChanged.emit()
        except Exception as e:
            print(e)

    def change_color(self, row, column, color):
        ix = self.index(row, column)
        self.colors[(row, column)] = color
        self.dataChanged.emit(ix, ix, (Qt.BackgroundRole,))


class TableViewer(QtWidgets.QMainWindow):

    def __init__(self):
        super(TableViewer, self).__init__()

        self.ui = loadUi("QTableViewForm.ui", self)
        self.ui.cmdRun1.clicked.connect(self.RunFunction1)
        self.showdata()

    def showdata(self):
        data = pd.read_pickle("productdata.pkl")
        self.model = PandasModel(data)
        self.ui.tableData.setModel(self.model)

    def set_cell_color(self, row, column):
        self.model.change_color(row, column, QBrush(Qt.red))

    def RunFunction1(self):
        start = timeit.default_timer()
        print("Start RunFunction1")

        print("Stop RunFunction1")
        end = timeit.default_timer()
        print("Process Time: ", (end - start))

if __name__ == "__main__":
    import sys
    app = QtWidgets.QApplication(sys.argv)
    win = TableViewer()
    win.set_cell_color(row=7,column=7)
    win.show()
    sys.exit(app.exec_())

When i run this i get cell (7,7) with a red background. Now i would like to do this by clicking the button and running the function Run Function1 but i can't get this to work.

Answer 1

Got it working , i just go over the row to get it in color. I'm working a lot with dataframes so if annyone has any suggestions to complete the PandasModel(QAbstractTableModel) class feel free ;-)

Updated Code:

import timeit
import pandas as pd
from PyQt5 import QtWidgets
from PyQt5.QtCore import *
from PyQt5.QtCore import QAbstractTableModel, Qt
from PyQt5.QtGui import *
from PyQt5.uic import loadUi


class PandasModel(QAbstractTableModel):
    def __init__(self, data, parent=None):
        QAbstractItemModel.__init__(self, parent)
        self._data = data
        self.colors = dict()

    def rowCount(self, parent=None):
        return self._data.index.size
    def columnCount(self, parent=None):
        return self._data.columns.size

    def setData(self, index, value, role):

        if role == Qt.EditRole:
            # Set the value into the frame.
            self._data.iloc[index.row(), index.column()] = value
            return True

    def data(self, index, role=Qt.DisplayRole):
        if index.isValid():
            if role == Qt.DisplayRole:
                return str(self._data.iloc[index.row(), index.column()])
            if role == Qt.EditRole:
                return str(self._data.iloc[index.row(), index.column()])
            if role == Qt.BackgroundRole:
                color = self.colors.get((index.row(), index.column()))
                if color is not None:
                    return color
        return None

    def headerData(self, rowcol, orientation, role):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self._data.columns[rowcol]
        if orientation == Qt.Vertical and role == Qt.DisplayRole:
            return self._data.index[rowcol]
        return None

    def change_color(self, row, column, color):
        ix = self.index(row, column)
        self.colors[(row, column)] = color
        self.dataChanged.emit(ix, ix, (Qt.BackgroundRole,))


class TableViewer(QtWidgets.QMainWindow):

    def __init__(self):
        super(TableViewer, self).__init__()

        self.ui = loadUi("QTableViewForm.ui", self)
        self.ui.cmdRun1.clicked.connect(self.RunFunction1)

        self.showdata()

    def showdata(self):
        data = pd.read_pickle("productdata.pkl")
        self.model = PandasModel(data)
        # self.ui.tableData.setModel(self.model)
        self.proxy_model = QSortFilterProxyModel()
        self.proxy_model.setFilterKeyColumn(-1)  # Search all columns.
        self.proxy_model.setSourceModel(self.model)
        self.proxy_model.sort(0, Qt.AscendingOrder)
        self.tableData.setModel(self.proxy_model)

    def set_cell_color(self, row, column):
        self.model.change_color(row, column, QBrush(Qt.red))

    def RunFunction1(self):
        start = timeit.default_timer()
        print("Start RunFunction1")
        #Gans de rij in 't rood
        colums = self.proxy_model.columnCount()

        for c in range(colums):

            self.set_cell_color(3, c)

        print("Stop RunFunction1")
        end = timeit.default_timer()
        print("Process Time: ", (end - start))


if __name__ == "__main__":
    import sys
    app = QtWidgets.QApplication(sys.argv)
    win = TableViewer()
    win.show()
    sys.exit(app.exec_())