如何为带日期索引的QAbstractTableModel插入指定格式新行？

我正在使用PyQt开发一款小型应用，用于展示数据库中的表格，支持用户添加新行并填写内容。表格的数据取自包含两列和日期索引的pd.DataFrame，要求插入的新行索引为最后一个日期索引加1天，且包含两个可编辑的空项。

当前应用无法实现新增行的需求，点击按钮后没有预期效果。以下是完整可复现代码：

import sys

import pandas as pd
from PyQt5.QtCore import QAbstractTableModel, QModelIndex, Qt
from PyQt5.QtWidgets import QApplication, QMainWindow, QTableView, QWidget, QGridLayout, QPushButton


class TableModel(QAbstractTableModel):

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

    def rowCount(self, parent=None) -> int:
        return self.table_data.shape[0]

    def columnCount(self, parent=None) -> int:
        return self.table_data.shape[1]

    def data(self, index: QModelIndex, role: int = ...):
        if role == Qt.DisplayRole:
            row = self.table_data.index[index.row()]
            col = self.table_data.columns[index.column()]
            value = self.table_data.loc[row][col]
            return str(value)

    def headerData(self, section: int, orientation: Qt.Orientation, role: int = ...):
        if role == Qt.DisplayRole:
            if orientation == Qt.Horizontal:
                return str(self.table_data.columns[section])
            if orientation == Qt.Vertical:
                return str(self.table_data.index[section])

    def insertRows(self, row: int, count: int, parent: QModelIndex = ...) -> bool:
        self.beginInsertRows()
        # whats here?
        self.endInsertRows()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = QMainWindow()

    widget = QWidget()

    layout = QGridLayout()
    widget.setLayout(layout)
    indexes = ['2023-01-03', '2023-01-02', '2023-01-01']
    data = {'col1': ['aa', 'bb', 'cc'], 'col2': [44, 55, 66]}
    table_data = pd.DataFrame(data=data, index=indexes)
    model = TableModel(table_data=table_data)
    table = QTableView()
    table.setModel(model)
    layout.addWidget(table)

    button = QPushButton("add row")
    button.clicked.connect(model.insertRows)

    button_2 = QPushButton("save changes")

    layout.addWidget(button)
    layout.addWidget(button_2)
    win.setCentralWidget(widget)

    win.show()
    app.exec()

要实现需求，需对代码做以下核心修改：

1. 实现表格编辑功能 #

在TableModel类中添加flags和setData方法，让单元格支持编辑，并将修改同步到DataFrame：

def flags(self, index: QModelIndex) -> Qt.ItemFlags:
    return Qt.ItemIsEditable | Qt.ItemIsSelectable | Qt.ItemIsEnabled

def setData(self, index: QModelIndex, value, role: int = ...) -> bool:
    if role == Qt.EditRole:
        row_idx = self.table_data.index[index.row()]
        col_name = self.table_data.columns[index.column()]
        self.table_data.loc[row_idx, col_name] = value
        self.dataChanged.emit(index, index)
        return True
    return False

2. 完善新增行逻辑 #

修改insertRows方法，计算新日期索引并添加空行，同时正确触发视图更新：

def insertRows(self, row: int, count: int, parent: QModelIndex = QModelIndex()) -> bool:
    # 默认在最后一行插入
    insert_pos = self.rowCount() if row == -1 else row
    # 转换索引为日期格式，计算新日期
    date_indexes = pd.to_datetime(self.table_data.index)
    new_date = date_indexes.max() + pd.Timedelta(days=1)
    new_date_str = new_date.strftime('%Y-%m-%d')

    self.beginInsertRows(parent, insert_pos, insert_pos + count - 1)
    # 批量添加空行
    for _ in range(count):
        self.table_data.loc[new_date_str] = [None, None]
        new_date += pd.Timedelta(days=1)
        new_date_str = new_date.strftime('%Y-%m-%d')
    self.endInsertRows()
    return True

3. 修复按钮连接问题 #

按钮点击时需传递正确的插入参数，修改连接代码：

button.clicked.connect(lambda: model.insertRows(-1, 1))

完整修改后的代码 #

import sys

import pandas as pd
from PyQt5.QtCore import QAbstractTableModel, QModelIndex, Qt
from PyQt5.QtWidgets import QApplication, QMainWindow, QTableView, QWidget, QGridLayout, QPushButton


class TableModel(QAbstractTableModel):

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

    def rowCount(self, parent=None) -> int:
        return self.table_data.shape[0]

    def columnCount(self, parent=None) -> int:
        return self.table_data.shape[1]

    def data(self, index: QModelIndex, role: int = ...):
        if role == Qt.DisplayRole or role == Qt.EditRole:
            row = self.table_data.index[index.row()]
            col = self.table_data.columns[index.column()]
            value = self.table_data.loc[row][col]
            return str(value) if value is not None else ""

    def headerData(self, section: int, orientation: Qt.Orientation, role: int = ...):
        if role == Qt.DisplayRole:
            if orientation == Qt.Horizontal:
                return str(self.table_data.columns[section])
            if orientation == Qt.Vertical:
                return str(self.table_data.index[section])

    def flags(self, index: QModelIndex) -> Qt.ItemFlags:
        return Qt.ItemIsEditable | Qt.ItemIsSelectable | Qt.ItemIsEnabled

    def setData(self, index: QModelIndex, value, role: int = ...) -> bool:
        if role == Qt.EditRole:
            row_idx = self.table_data.index[index.row()]
            col_name = self.table_data.columns[index.column()]
            self.table_data.loc[row_idx, col_name] = value
            self.dataChanged.emit(index, index)
            return True
        return False

    def insertRows(self, row: int, count: int, parent: QModelIndex = QModelIndex()) -> bool:
        insert_pos = self.rowCount() if row == -1 else row
        date_indexes = pd.to_datetime(self.table_data.index)
        new_date = date_indexes.max() + pd.Timedelta(days=1)
        new_date_str = new_date.strftime('%Y-%m-%d')

        self.beginInsertRows(parent, insert_pos, insert_pos + count - 1)
        for _ in range(count):
            self.table_data.loc[new_date_str] = [None, None]
            new_date += pd.Timedelta(days=1)
            new_date_str = new_date.strftime('%Y-%m-%d')
        self.endInsertRows()
        return True


if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = QMainWindow()

    widget = QWidget()

    layout = QGridLayout()
    widget.setLayout(layout)
    indexes = ['2023-01-03', '2023-01-02', '2023-01-01']
    data = {'col1': ['aa', 'bb', 'cc'], 'col2': [44, 55, 66]}
    table_data = pd.DataFrame(data=data, index=indexes)
    model = TableModel(table_data=table_data)
    table = QTableView()
    table.setModel(model)
    layout.addWidget(table)

    button = QPushButton("add row")
    button.clicked.connect(lambda: model.insertRows(-1, 1))

    button_2 = QPushButton("save changes")

    layout.addWidget(button)
    layout.addWidget(button_2)
    win.setCentralWidget(widget)

    win.show()
    app.exec()

关键说明 #

• flags方法设置单元格的可编辑、可选择属性，是编辑功能的基础。
• setData方法将用户编辑的内容同步到DataFrame，并触发视图更新。
• insertRows方法中，先将原索引转为日期格式，计算出最大日期加1天作为新行索引，再添加空行并通知视图刷新。
• 按钮通过lambda传递参数，指定在表格末尾插入1条新行。

内容的提问来源于stack exchange，提问作者papadulo