Generic QT model with undo integration, and batch editing

generic_table_model.py

"""
A collection of classes for working with MVC in QT, including a generic table model, and an example of integrating undo via a proxy model.
"""
import typing
from qtpy import QtCore, QtGui, QtWidgets


T = typing.TypeVar("T")

class SchemaColumn(typing.Generic[T]):
    """
    SchemaColumn is an abstract base class that defines the interface for a column in the SchemaTableModel.
    It provides methods to get the header name, determine the item flags for a given item, and get/set data for a given item and role.
    """
    def header(self) -> str:
        raise NotImplementedError()

    def flags(self, item:T) -> QtCore.Qt.ItemFlag:
        return QtCore.Qt.ItemFlag.ItemIsEnabled | QtCore.Qt.ItemFlag.ItemIsSelectable
    
    def getData(self, item: T, role:QtCore.Qt.ItemDataRole) -> typing.Any:
        raise NotImplementedError()
    
    def setData(self, item: T, value: typing.Any, role:QtCore.Qt.ItemDataRole) -> bool:
        raise NotImplementedError()
    

class SimpleObjectColumn(SchemaColumn[T]):
    """
    A simple implementation of SchemaColumn that uses getter and setter callables to access the data for a given item.
    """
    def __init__(self, header: str, getter: typing.Callable[[T], typing.Any], setter: typing.Optional[typing.Callable[[T, typing.Any], None]] = None):
        self._header = header
        self._getter = getter
        self._setter = setter
    
    def header(self) -> str:
        return self._header

    def flags(self, item: T) -> QtCore.Qt.ItemFlag:
        flags = super().flags(item)
        if self._setter is not None:
            flags |= QtCore.Qt.ItemFlag.ItemIsEditable
        return flags

    
    def getData(self, item: T, role: QtCore.Qt.ItemDataRole) -> typing.Any:
        if role == QtCore.Qt.ItemDataRole.DisplayRole:
            return self._getter(item)
        elif role == QtCore.Qt.ItemDataRole.EditRole:
            return self._getter(item)
        return None
    
    def setData(self, item: T, value: typing.Any, role: QtCore.Qt.ItemDataRole) -> bool:
        if self._setter is not None and role == QtCore.Qt.ItemDataRole.EditRole:
            print(f"Setting data: {value} on item: {item} type: {type(value)}")
            self._setter(item, value)
            return True
        return False

    
class SchemaTableModel(QtCore.QAbstractTableModel, typing.Generic[T]):
    """
    A generic model for displaying a list of items of type T in a table view, where the columns are defined by a list of SchemaColumn[T] instances.

    Example:
        class TestItem:
            def __init__(self, name: str, value: int):
                self.name = name
                self.value = value

        columns = [
            SimpleObjectColumn[TestItem]("Name", lambda item: item.name, lambda item, value: setattr(item, "name", value)),
            SimpleObjectColumn[TestItem]("Value", lambda item: item.value, lambda item, value: setattr(item, "value", int(value or 0))),
        ]

        model = SchemaTableModel[TestItem](columns)
        model._items = [TestItem("Item 1", 10), TestItem("Item 2", 20)]
        view = QtWidgets.QTableView()
        view.setModel(model)

    """
    def __init__(self, columns: typing.Sequence[SchemaColumn[T]], parent: QtCore.QObject | None = None):
        super().__init__(parent)
        self._columns = columns
        self._items: list[T] = []

    def rowCount(self, parent: QtCore.QModelIndex) -> int:
        return len(self._items)
    
    def columnCount(self, parent: QtCore.QModelIndex) -> int:
        return len(self._columns)
    
    def headerData(self, section: int, orientation: QtCore.Qt.Orientation, role: QtCore.Qt.ItemDataRole) -> typing.Any:
        if role == QtCore.Qt.ItemDataRole.DisplayRole and orientation == QtCore.Qt.Orientation.Horizontal:
            return self._columns[section].header()
        
        return super().headerData(section, orientation, role)
    
    def itemFromIndex(self, index: QtCore.QModelIndex) -> T | None:
        if not index.isValid():
            return None
        return self._items[index.row()]
    
    def flags(self, index: QtCore.QModelIndex) -> QtCore.Qt.ItemFlag:
        if not index.isValid():
            return QtCore.Qt.ItemFlag.NoItemFlags
        
        item = self._items[index.row()]
        column = self._columns[index.column()]
        return column.flags(item)
    
    def data(self, index: QtCore.QModelIndex, role: QtCore.Qt.ItemDataRole) -> typing.Any:
        if not index.isValid():
            return None
        
        item = self._items[index.row()]
        column = self._columns[index.column()]
        return column.getData(item, role)
    
    def setData(self, index: QtCore.QModelIndex, value: typing.Any, role: QtCore.Qt.ItemDataRole) -> bool:
        if not index.isValid():
            return False
        
        item = self._items[index.row()]
        column = self._columns[index.column()]
        if column.setData(item, value, role):
            self.dataChanged.emit(index, index, [role])
            return True
        return False



class ChangeValueCommand(QtGui.QUndoCommand):
    """
    A QUndoCommand that changes a value in the model. It stores the model, index, new value and old value, 
    and applies the change in redo() and reverts it in undo().
    """
    def __init__(self, model: QtCore.QAbstractItemModel, index: QtCore.QModelIndex, new_value: typing.Any, old_value: typing.Any, parent=None):
        super().__init__(parent=parent)
        self._model = model
        self._index = index
        self._new_value = new_value
        self._old_value = old_value
        self._dont_apply = True


    def redo(self):
        # If the change has already been applied to the model, we don't want to apply it again when redo is called for the first time.
        if self._dont_apply:
            self._dont_apply = False
            return
        self._model.setData(self._index, self._new_value, QtCore.Qt.ItemDataRole.EditRole)

    def undo(self):
        self._model.setData(self._index, self._old_value, QtCore.Qt.ItemDataRole.EditRole)


class UndoProxyModel(QtCore.QIdentityProxyModel):
    """
    A proxy model that adds undo/redo functionality to a source model by intercepting setData calls.
    """
    def __init__(self, undo_stack: QtGui.QUndoStack, parent: QtCore.QObject | None = None):
        super().__init__(parent)
        self._undo_stack = undo_stack

    def setData(self, index: QtCore.QModelIndex, value: typing.Any, role: QtCore.Qt.ItemDataRole) -> bool:
        if not index.isValid():
            return False
        
        source_index = self.mapToSource(index)
        old_value = self.sourceModel().data(source_index, role)
        if old_value == value:
            return False
        
        success = self.sourceModel().setData(source_index, value, role)
        if success:
            # Note: the first invokation of redo() is skipped, since the change has already been applied to the model by setData()
            command = ChangeValueCommand(self.sourceModel(), source_index, value, old_value)
            self._undo_stack.push(command)        
        
        return True




class BatchEditingBehavior(QtCore.QObject):
    """
    A behavior that can be attached to a QTableView to enable batch editing of a column for multiple selected rows.

    edit begin/end is signaled to allow for undo stack macros to be used for batch edits.

    Attibutes:
        editBatchBegin: A signal that is emitted when a batch edit is initiated, with the column header as an argument.
        editBatchEnd: A signal that is emitted when a batch edit is completed.
        _edit_role: The item data role that should be used for editing (default is QtCore.Qt.ItemDataRole.EditRole)
      
    """
    editBatchBegin = QtCore.Signal(str) # type: ignore
    editBatchEnd = QtCore.Signal() # type: ignore

    def __init__(self, edit_role: QtCore.Qt.ItemDataRole = QtCore.Qt.ItemDataRole.EditRole, parent: QtCore.QObject | None = None):
        super().__init__(parent)
        self._edit_role = edit_role


    def contextMenuRequested(self, pos: QtCore.QPoint):
        """
        Slot for handling the customContextMenuRequested signal from the view.
        """
        view = self.sender()
        if not isinstance(view, QtWidgets.QTableView):
            return
        actions = self.getActions(view, pos)
        if not actions:
            return
        
        menu = QtWidgets.QMenu(view)
        menu.exec(actions, pos)

    def getActions(self, view: QtWidgets.QTableView, pos: QtCore.QPoint) -> typing.Sequence[QtWidgets.QAction]:
        """
        Generate a list of actions for the context menu based on the current selection in the view.
        This can potentiall be overridden in a subclass to extend the available actions.
        """
        result: typing.List[QtWidgets.QAction] = []
        index = view.indexAt(pos)

        if not index.isValid():
            return result
        
        model = view.model()
        if not isinstance(model, QtCore.QAbstractItemModel):
            return result
        
        selected_indexes = view.selectionModel().selectedIndexes()
        if len(selected_indexes) <= 1:
            return result
        
        selected_indexes.remove(index)

        header = model.headerData(index.column(), QtCore.Qt.Orientation.Horizontal, QtCore.Qt.ItemDataRole.DisplayRole)

        edit_action = QtWidgets.QAction(f"Batch Edit: {header}", view) # type: ignore
        edit_action.triggered.connect(self.onBatchEdit)

        edit_action.setProperty("view", view)
        edit_action.setProperty("index", index)
        edit_action.setProperty("other_indexes", selected_indexes)
        edit_action.setProperty("header", header)
        result.append(edit_action)
        return result      
    

    def onBatchEdit(self):
        """
        Slot for handling the triggered signal from the batch edit action. 
        """
        action: QtWidgets.QAction = self.sender()  # type: ignore
        view: QtWidgets.QTableView = action.property("view")
        index: QtCore.QModelIndex = action.property("index")        
        other_indexes: typing.List[QtCore.QModelIndex] = action.property("other_indexes")
        header: str = action.property("header")

        model = view.model()
        if not isinstance(model, QtCore.QAbstractItemModel):
            return
        
        delegate = view.itemDelegate(index)
        if not isinstance(delegate, QtWidgets.QStyledItemDelegate):
            return
        
        self.editBatchBegin.emit(header)
        
        def onCommit():
            new_value = model.data(index, QtCore.Qt.ItemDataRole.EditRole)
            for other_index in other_indexes:
                model.setData(other_index, new_value, self._edit_role)

            delegate.commitData.disconnect(onCommit)

        def onClose():
            delegate.commitData.disconnect(onCommit)
            delegate.closeEditor.disconnect(onClose)
            self.editBatchEnd.emit()

        delegate.commitData.connect(onCommit)
        delegate.closeEditor.connect(onClose)

        view.edit(index)
        



        



if __name__ == "__main__":
    import sys

    class TestItem:
        def __init__(self, name: str, value: int, metadata: dict[str, typing.Any]):
            self.name = name
            self.value = value
            self.metadata = metadata

    app = QtWidgets.QApplication(sys.argv)

    # define how to display and edit the TestItem in the table model using SimpleObjectColumn
    columns = [
        SimpleObjectColumn[TestItem]("Name", lambda item: item.name, lambda item, value: setattr(item, "name", value)),
        SimpleObjectColumn[TestItem]("Value", lambda item: item.value, lambda item, value: setattr(item, "value", int(value or 0))),
        SimpleObjectColumn[TestItem]("Metadata", lambda item: item.metadata.get("type"), lambda item, value: item.metadata.__setitem__("type", value)),
    ]

    model = SchemaTableModel[TestItem](columns)

    # add some test items to the model
    model._items = [TestItem("Item 1", 10, {"type": "A"}), TestItem("Item 2", 20, {"type": "B"})]
    
    # set up the undo stack and proxy model
    undo_stack = QtGui.QUndoStack()
    undo_proxy_model = UndoProxyModel(undo_stack)
    undo_proxy_model.setSourceModel(model)

    undo_action = undo_stack.createUndoAction(app, "Undo")
    undo_action.setShortcut(QtGui.QKeySequence.StandardKey.Undo)

    redo_action = undo_stack.createRedoAction(app, "Redo")
    redo_action.setShortcut(QtGui.QKeySequence.StandardKey.Redo)
    

    view = QtWidgets.QTableView()
    view.setModel(undo_proxy_model)
    view.setContextMenuPolicy(QtCore.Qt.ContextMenuPolicy.CustomContextMenu)
    view.addAction(undo_action)
    view.addAction(redo_action)
    view.show()

    behavior = BatchEditingBehavior()

    def beginBatchEdit(header: str):
        print(f"Beginning batch edit for column: {header}")
        undo_stack.beginMacro(f"Batch Edit: {header}")

    def endBatchEdit():
        print("Ending batch edit")
        undo_stack.endMacro()

    behavior.editBatchBegin.connect(beginBatchEdit)
    behavior.editBatchEnd.connect(endBatchEdit)
 
    view.customContextMenuRequested.connect(behavior.contextMenuRequested)



    sys.exit(app.exec())