Пример #1
0
class EditableTreeView(QWidget):
    dataChanged = Signal()
    selectionChanged = Signal()

    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.__stack: List = []
        self.__stack_index: int = -1

        def push_on_data_changed(_, __, roles):
            if Qt.EditRole in roles:
                self._push_data()

        self.__model = QStandardItemModel()
        self.__model.dataChanged.connect(self.dataChanged)
        self.__model.dataChanged.connect(push_on_data_changed)
        self.__root: QStandardItem = self.__model.invisibleRootItem()

        self.__tree = TreeView(self.dataChanged)
        self.__tree.drop_finished.connect(self.dataChanged)
        self.__tree.drop_finished.connect(self._push_data)
        self.__tree.setModel(self.__model)
        self.__tree.selectionModel().selectionChanged.connect(
            self.selectionChanged)

        actions_widget = ModelActionsWidget()
        actions_widget.layout().setSpacing(1)

        action = QAction("+", self, toolTip="Add a new word")
        action.triggered.connect(self.__on_add)
        actions_widget.addAction(action)

        action = QAction("\N{MINUS SIGN}", self, toolTip="Remove word")
        action.triggered.connect(self.__on_remove)
        actions_widget.addAction(action)

        action = QAction("\N{MINUS SIGN}R", self,
                         toolTip="Remove word recursively (incl. children)")
        action.triggered.connect(self.__on_remove_recursive)
        actions_widget.addAction(action)

        gui.rubber(actions_widget)

        self.__undo_action = action = QAction("Undo", self, toolTip="Undo")
        action.triggered.connect(self.__on_undo)
        actions_widget.addAction(action)

        self.__redo_action = action = QAction("Redo", self, toolTip="Redo")
        action.triggered.connect(self.__on_redo)
        actions_widget.addAction(action)

        self._enable_undo_redo()

        layout = QVBoxLayout()
        layout.setSpacing(1)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.addWidget(self.__tree)
        layout.addWidget(actions_widget)
        self.setLayout(layout)

    def __on_add(self):
        parent: QStandardItem = self.__root
        selection: List = self.__tree.selectionModel().selectedIndexes()
        if selection:
            sel_index: QModelIndex = selection[0]
            parent: QStandardItem = self.__model.itemFromIndex(sel_index)

        item = QStandardItem("")
        parent.appendRow(item)
        index: QModelIndex = item.index()
        with disconnected(self.__model.dataChanged, self.dataChanged):
            self.__model.setItemData(index, {Qt.EditRole: ""})
        self.__tree.setCurrentIndex(index)
        self.__tree.edit(index)

    def __on_remove_recursive(self):
        sel_model: QItemSelectionModel = self.__tree.selectionModel()
        if len(sel_model.selectedIndexes()):
            while sel_model.selectedIndexes():
                index: QModelIndex = sel_model.selectedIndexes()[0]
                self.__model.removeRow(index.row(), index.parent())
            self._push_data()
            self.dataChanged.emit()

    def __on_remove(self):
        sel_model: QItemSelectionModel = self.__tree.selectionModel()
        if len(sel_model.selectedIndexes()):

            while sel_model.selectedIndexes():
                index: QModelIndex = sel_model.selectedIndexes()[0]

                # move children to item's parent
                item: QStandardItem = self.__model.itemFromIndex(index)
                children = [item.takeChild(i) for i in range(item.rowCount())]
                parent = item.parent() or self.__root

                self.__model.removeRow(index.row(), index.parent())

                for child in children[::-1]:
                    parent.insertRow(index.row(), child)

            self.__tree.expandAll()
            self._push_data()
            self.dataChanged.emit()

    def __on_undo(self):
        self.__stack_index -= 1
        self._set_from_stack()

    def __on_redo(self):
        self.__stack_index += 1
        self._set_from_stack()

    def get_words(self) -> List:
        return _model_to_words(self.__root)

    def get_selected_words(self) -> Set:
        return set(self.__model.itemFromIndex(index).text() for index in
                   self.__tree.selectionModel().selectedIndexes())

    def get_selected_words_with_children(self) -> Set:
        words = set()
        for index in self.__tree.selectionModel().selectedIndexes():
            item: QStandardItem = self.__model.itemFromIndex(index)
            words.update(_model_to_words(item))
        return words

    def get_data(self, with_selection=False) -> Union[Dict, OntoType]:
        selection = self.__tree.selectionModel().selectedIndexes()
        return _model_to_tree(self.__root, selection, with_selection)

    def set_data(self, data: Dict, keep_history: bool = False):
        if not keep_history:
            self.__stack = []
            self.__stack_index = -1
        self._set_data(data)
        self._push_data()

    def _set_data(self, data: Dict):
        self.clear()
        _tree_to_model(data, self.__root, self.__tree.selectionModel())
        self.__tree.expandAll()

    def clear(self):
        if self.__model.hasChildren():
            self.__model.removeRows(0, self.__model.rowCount())

    def _enable_undo_redo(self):
        index = self.__stack_index
        self.__undo_action.setEnabled(index >= 1)
        self.__redo_action.setEnabled(index < len(self.__stack) - 1)

    def _push_data(self):
        self.__stack_index += 1
        self.__stack = self.__stack[:self.__stack_index]
        self.__stack.append(self.get_data())
        self._enable_undo_redo()

    def _set_from_stack(self):
        assert self.__stack_index < len(self.__stack)
        assert self.__stack_index >= 0
        self._set_data(self.__stack[self.__stack_index])
        self._enable_undo_redo()
        self.dataChanged.emit()