PyQt5 QTableViewにpandasを表示するためのQAbstractItemModelを作成する

PyQt5でQTableViewにpandasのDataFrameを表示するためのQAbstractItemModelを作ってみました。

表示するデータの取得

QAbstractItemModelのデータを返す関数はdef data(self, index, role=QtCore.Qt.DisplayRole)です。この関数はindexで取得するデータの位置を指定して、その位置のデータを返します。そこでindexの行、列番号をindex.row()とindex.column()で取得して、それらの番号でDataFrameの特定の場所のデータをDataFrame.iat[row, column]で取得して、data関数の返り値とします。

def data(self, index, role=QtCore.Qt.DisplayRole):
    if not index.isValid():
        return QtCore.QVariant()
    if role == QtCore.Qt.EditRole or role == QtCore.Qt.DisplayRole:
        return self.df.iat[index.row(), index.column()]
    return QtCore.QVariant()

データのセット

QAbstractItemModelのデータを返す関数はdef setData(self, index, value, role=QtCore.Qt.EditRole)です。この関数はindexで取得するデータの位置を指定して、その位置にデータvalueをセットします。データのセットにはデータの表示と同じDataFrame.iatを使用します。

def setData(self, index, value, role=QtCore.Qt.EditRole):
    if role == QtCore.Qt.EditRole:
        self.df.iat[index.row(), index.column()] = value
        return True
    return False

ヘッダーの表示

ヘッダーの方向が縦か横のどちらかを示すorientationによってヘッドーの表示内容を変えます。縦の場合はDataFrameのindexを表示して、横の場合はDataFrameのcolumnsを表示します。

def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
    if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
        return self.df.columns[section]
    if orientation == QtCore.Qt.Vertical and role == QtCore.Qt.DisplayRole:
        return self.df.index[section]

表のサイズ

表の縦横のサイズはそれぞれ、def columnCount(self, parent=QtCore.QModelIndex())とdef rowCount(self, parent=QtCore.QModelIndex())という関数です。これらはDataFrame.shapeから取得した値をそのまま返す関数としました。

def columnCount(self, parent=QtCore.QModelIndex()):
    return self.df.shape[1]
def rowCount(self, parent=QtCore.QModelIndex()):
    return self.df.shape[0]

列の挿入

列の挿入にはdef insertColumns(self, column, count, parent=QtCore.QModelIndex())を使用します。挿入を行う前にbeginInsertColumnsを実行して、挿入完了後にendInsertColumnsを実行する必要があります。挿入完了前にendInsertColumnsを実行したりするとエラーがが発生します。挿入位置はcolumnからcountの個数挿入します。そこでDataFrame.ilocで列を指定してcolumnの位置でDataFrameを左側と右側にスライスします。その中間に挿入されるDataFrameを、indexを既存のデータフレームから流用して、columnsを挿入する個数分の数として作成します。そして、pandas.concatでスライスした左側、作成した中間、スライスした右側を連結して新しいDataFrameを作成することで列の挿入を実装します。

def insertColumns(self, column, count, parent=QtCore.QModelIndex()):
    self.beginInsertColumns(parent, column, column + count - 1)
    columns = [ str(self.columnCount()+i+1) for i in range(count) ]
    left = self.df.iloc[:, 0:column]
    mid = pd.DataFrame(index=self.df.index, columns=columns)
    right = self.df.iloc[:, column+count-1:self.columnCount()]
    self.df = pd.concat( [left, mid, right], axis=1 )
    
    self.endInsertColumns()

行の挿入

行の挿入も列の挿入と同じ要領です。挿入位置rowでDataFrameをスライスして、新たに作成したDataFrameとそれらを連結して新しいDataFrameを作成することで行の挿入を実装します。

def insertRows(self, row, count, parent=QtCore.QModelIndex()):
    self.beginInsertRows(parent, row, row + count - 1)
    indexes = [ str(self.rowCount()+i) for i in range(count) ]
    left = self.df[0:row]
    mid = pd.DataFrame(index=indexes, columns=self.df.columns)
    right = self.df[row+count-1:self.rowCount()]
    self.df = pd.concat([left, mid, right])
    self.endInsertRows()

列の削除

列の削除は、列の挿入の処理から中間に挿入するDataFrameを作成する処理を除いた処理です。

def removeColumns(self, column, count, parent=QtCore.QModelIndex()):
    self.beginRemoveColumns(parent, column, column + count - 1)
    left = self.df.iloc[:, 0:column]
    right = self.df.iloc[:, column+count:self.columnCount()]
    self.df = pd.concat( [left, right], axis=1 )
    self.endRemoveColumns()

行の削除

行の削除も列の削除と同様の処理です。

def removeRows(self, row, count, parent=QtCore.QModelIndex()):
    self.beginRemoveRows(parent, row, row + count - 1)
    
    left = self.df.iloc[0:row]
    right = self.df.iloc[row+count:self.rowCount()]
    self.df = pd.concat( [left, right], axis=0 )
    self.endRemoveRows()

ソースコード

import pandas as pd
from PyQt5 import QtWidgets, QtCore
class PandasModel(QtCore.QAbstractItemModel):
    def __init__(self, parent=None, data_frame=None):
        super(PandasModel, self).__init__(parent)
        self.df = data_frame
    def columnCount(self, parent=QtCore.QModelIndex()):
        return self.df.shape[1]
    def data(self, index, role=QtCore.Qt.DisplayRole):
        if not index.isValid():
            return QtCore.QVariant()
        if role == QtCore.Qt.EditRole or role == QtCore.Qt.DisplayRole:
            return self.df.iat[index.row(), index.column()]
        return QtCore.QVariant()
    def flags(self, index):
        return QtCore.Qt.ItemIsEditable | QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            return self.df.columns[section]
        if orientation == QtCore.Qt.Vertical and role == QtCore.Qt.DisplayRole:
            return self.df.index[section]
    def index(self, row, column, parent=QtCore.QModelIndex()):
        if not self.hasIndex(row, column, parent):
            return QtCore.QModelIndex()
        return self.createIndex(row, column, QtCore.QModelIndex())
    def insertColumns(self, column, count, parent=QtCore.QModelIndex()):
        self.beginInsertColumns(parent, column, column + count - 1)
        columns = [ str(self.columnCount()+i+1) for i in range(count) ]
        left = self.df.iloc[:, 0:column]
        mid = pd.DataFrame(index=self.df.index, columns=columns)
        right = self.df.iloc[:, column+count-1:self.columnCount()]
        self.df = pd.concat( [left, mid, right], axis=1 )
        
        self.endInsertColumns()
    def insertRows(self, row, count, parent=QtCore.QModelIndex()):
        self.beginInsertRows(parent, row, row + count - 1)
        indexes = [ str(self.rowCount()+i) for i in range(count) ]
        left = self.df[0:row]
        mid = pd.DataFrame(index=indexes, columns=self.df.columns)
        right = self.df[row+count-1:self.rowCount()]
        self.df = pd.concat([left, mid, right])
        self.endInsertRows()
    def removeColumns(self, column, count, parent=QtCore.QModelIndex()):
        self.beginRemoveColumns(parent, column, column + count - 1)
        left = self.df.iloc[:, 0:column]
        right = self.df.iloc[:, column+count:self.columnCount()]
        self.df = pd.concat( [left, right], axis=1 )
        self.endRemoveColumns()
    def removeRows(self, row, count, parent=QtCore.QModelIndex()):
        self.beginRemoveRows(parent, row, row + count - 1)
        
        left = self.df.iloc[0:row]
        right = self.df.iloc[row+count:self.rowCount()]
        self.df = pd.concat( [left, right], axis=0 )
        self.endRemoveRows()
        
    def rowCount(self, parent=QtCore.QModelIndex()):
        return self.df.shape[0]
    def setData(self, index, value, role=QtCore.Qt.EditRole):
        if role == QtCore.Qt.EditRole:
            self.df.iat[index.row(), index.column()] = value
            return True
        return False
class Delegate(QtWidgets.QStyledItemDelegate):
    def __init__(self, parent=None, setModelDataEvent=None):
        super(Delegate, self).__init__(parent)
        self.setModelDataEvent = setModelDataEvent
 
    def createEditor(self, parent, option, index):
        index.model().data(index, QtCore.Qt.DisplayRole)
        return QtWidgets.QLineEdit(parent)
 
    def setEditorData(self, editor, index):
        value = index.model().data(index, QtCore.Qt.DisplayRole)
        editor.setText(str(value))
        
    def setModelData(self, editor, model, index):
        model.setData(index, editor.text())
        if not self.setModelDataEvent is None:
            self.setModelDataEvent()
def main():
    
    def insert_row():
        if len(tableView.selectedIndexes()) == 0:
            model.insertRows(model.rowCount(), 1)
        else:
            rows = list( set([ index.row() for index in tableView.selectedIndexes() ]) )
            for row in rows[::-1]:
                model.insertRows(row, 1)
    def delete_row():
        rows = list( set([ index.row() for index in tableView.selectedIndexes() ]) )
        for row in rows[::-1]:
            model.removeRows(row, 1)
    def insert_column():
        if len(tableView.selectedIndexes()) == 0:
            model.insertColumns(model.columnCount(), 1)
        else:
            columns = list( set([ index.column() for index in tableView.selectedIndexes() ]) )
            for column in columns[::-1]:
                model.insertColumns(column, 1)
    def delete_column():
        columns = list( set([ index.column() for index in tableView.selectedIndexes() ]) )
        for column in columns[::-1]:
            model.removeColumns(column, 1)
    def exec_context_menu(point):
        menu = QtWidgets.QMenu(tableView)
        menu.addAction('Insert row', insert_row)
        menu.addAction('Delete row', delete_row)
        menu.addAction('Insert column', insert_column)
        menu.addAction('Delete column', delete_column)
        menu.exec( tableView.focusWidget().mapToGlobal(point) )
    
    import sys
    app = QtWidgets.QApplication(sys.argv)
    window = QtWidgets.QMainWindow()
    window.resize(600, 600)
    centralwidget = QtWidgets.QWidget(window)
    verticalLayout = QtWidgets.QVBoxLayout(centralwidget)
    df = pd.DataFrame(
        {
            '1': ['A1', 'A2', 'A3'],
            '2': ['B1', 'B2', 'B3'],
            '3': ['C1', 'C2', 'C3']
        }, index=['0', '1', '2']
    )
    
    tableView = QtWidgets.QTableView(centralwidget)
    model = PandasModel(tableView, df)
    tableView.setModel(model)
    tableView.setItemDelegate(Delegate())
    tableView.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
    tableView.customContextMenuRequested.connect(exec_context_menu)
    verticalLayout.addWidget(tableView)
    window.setCentralWidget(centralwidget)
    window.show()
    app.exec()
if __name__ == '__main__':
    main()