from collections.abc import Callable from functools import partial from typing import Protocol from PySide6.QtCore import QObject, QPoint, QSize, Qt from PySide6.QtWidgets import QAbstractItemView, QDialog, QHeaderView, QMenu from bedit_core.models import Component, ComponentID, EquationImplementation, GraphImplementation, Port, PortID, Parameter, ParameterID from bedit_core.models import Document as CoreDocument from bedit_gui.documents import Document from bedit_gui.models import Icon from bedit_gui.views.dialogs.interface_editor_dialog import InterfaceEditorDialog from bedit_gui.views.dialogs.param_editor_dialog import ParamEditorDialog from bedit_gui.views.icon_editor_window import IconEditorWindow from bedit_gui.views.main_window import MainWindow from bedit_gui.views.models.document_tree_model import DocumentTreeModel from bedit_gui.utils.icon import render_icon ICON_SIZE = QSize(32, 32) class InterfaceEditorLike(Protocol): def exec(self) -> int: ... def ports(self) -> dict[PortID, Port]: ... class ParamEditorLike(Protocol): def exec(self) -> int: ... def params(self) -> dict[ParameterID, Parameter]: ... InterfaceEditorFactory = Callable[ [dict[PortID, Port], MainWindow], InterfaceEditorLike, ] ParamEditorFactory = Callable[ [dict[ParameterID, Parameter], MainWindow], ParamEditorLike ] class DocumentTreeController(QObject): def __init__( self, document: Document, window: MainWindow, interface_editor_factory: InterfaceEditorFactory = InterfaceEditorDialog, param_editor_factory: ParamEditorFactory = ParamEditorDialog, ) -> None: super().__init__(window) self.document = document self.window = window self.model = DocumentTreeModel() self.interface_editor_factory = interface_editor_factory self.param_editor_factory = param_editor_factory self._icon_editors: list[IconEditorWindow] = [] self._components: dict[ComponentID, Component] = {} window.ui.documentTree.setModel(self.model) window.ui.documentTree.selectionModel().selectionChanged.connect(self._selection_changed) document.model_changed.connect(self._on_document_changed) document.icon_changed.connect(self._on_icon_changed) self.model.rename_document_requested.connect(self.document.rename) self.model.rename_component_requested.connect(self.document.rename_component) window.ui.actionDelete.triggered.connect(self.delete_selected_component) # Add deselection with esc to this widget window.ui.actionEscape.setShortcutContext(Qt.ShortcutContext.WidgetWithChildrenShortcut) window.ui.documentTree.addAction(window.ui.actionEscape) window.ui.actionEscape.triggered.connect(self.deselect) window.ui.documentTree.setHeaderHidden(True) window.ui.documentTree.setSelectionMode(QAbstractItemView.SelectionMode.ExtendedSelection) window.ui.documentTree.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows) window.ui.documentTree.setIconSize(QSize(48, 48)) window.ui.documentTree.header().setStretchLastSection(False) window.ui.documentTree.header().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch) window.ui.documentTree.header().setSectionResizeMode(1, QHeaderView.ResizeMode.Fixed) window.ui.documentTree.setColumnWidth(1, 56) window.ui.documentTree.setContextMenuPolicy( Qt.ContextMenuPolicy.CustomContextMenu ) window.ui.documentTree.customContextMenuRequested.connect( self._show_context_menu ) self._on_document_changed(document.model) def _on_document_changed(self, model: CoreDocument) -> None: """Rebuild the tree whenever New/Open replaces the core document.""" self._show_equation_component(None) self.model.set_document(model) self._components = {} self._collect_components(model.root) for component_id, component in self._components.items(): icon = self.document.component_icon(component_id) self.model.set_component_icon(component_id, render_icon(icon, component.interface.ports, ICON_SIZE)) # Optional presentation behavior. Later, you could instead remember # expanded component IDs and restore only those nodes. self.window.ui.documentTree.expandAll() def _selection_changed(self, *_args: object) -> None: indexes = self.window.ui.documentTree.selectionModel().selectedRows(0) component = self.model.value(indexes[0]) if len(indexes) == 1 else None self._show_equation_component(component if isinstance(component, Component) else None) def _show_equation_component(self, component: Component | None) -> None: if component is not None and isinstance(component.implementation, EquationImplementation): self.window.equation_editor.set_component(component) self.window.equation_editor.show() else: self.window.equation_editor.set_component(None) self.window.equation_editor.hide() def _on_icon_changed(self, component_id: ComponentID, icon: object) -> None: component = self._components.get(component_id) if component is None: return self.model.set_component_icon(component_id, render_icon(icon if isinstance(icon, Icon) else Icon(), component.interface.ports, ICON_SIZE)) def _collect_components(self, components: dict[ComponentID, Component]) -> None: for component_id, component in components.items(): self._components[component_id] = component if isinstance(component.implementation, GraphImplementation): self._collect_components(component.implementation.graph.components) def _show_context_menu(self, position: QPoint) -> None: index = self.window.ui.documentTree.indexAt(position) component = self.model.value(index) if not isinstance(component, Component): return menu = QMenu(self.window.ui.documentTree) edit_interface = menu.addAction("Edit Interface") edit_params = menu.addAction("Edit Parameters") edit_icon = menu.addAction("Edit Icon") menu.addSeparator() add_graph_component = menu.addAction("Add Graph Component") add_equation_component = menu.addAction("Add Equation Component") menu.addSeparator() delete_component = menu.addAction("Delete Component") selected = menu.exec( self.window.ui.documentTree.viewport().mapToGlobal(position) ) if selected is edit_interface: self._edit_interface(component) elif selected is edit_params: self._edit_params(component) elif selected is edit_icon: self._edit_icon(component) elif selected is add_graph_component: self._add_graph_component(component) elif selected is add_equation_component: self._add_equation_component(component) elif selected is delete_component: self._delete_component(component) def _edit_interface(self, component: Component) -> None: dialog = self.interface_editor_factory( component.interface.ports, self.window, ) if dialog.exec() == QDialog.DialogCode.Accepted: self.document.update_component_ports(component, dialog.ports()) def _edit_params(self, component: Component) -> None: dialog = self.param_editor_factory( component.parameters, self.window ) if dialog.exec() == QDialog.DialogCode.Accepted: self.document.update_component_params(component, dialog.params()) def _edit_icon(self, component: Component) -> None: component_id = self.document.component_id(component) editor = IconEditorWindow(self.document.component_icon(component_id), component.interface.ports, self.window) editor.saved.connect(partial(self.document.change_icon, component_id)) editor.destroyed.connect(partial(self._icon_editor_closed, editor)) self._icon_editors.append(editor) editor.show() def _icon_editor_closed(self, editor: IconEditorWindow, *_args: object) -> None: if editor in self._icon_editors: self._icon_editors.remove(editor) def _add_graph_component(self, component: Component) -> None: self.document.add_empty_graph_component(component) def _add_equation_component(self, component: Component) -> None: self.document.add_empty_equation_component(component) def deselect(self) -> None: self.window.ui.documentTree.selectionModel().clear() def delete_selected_component(self) -> None: focused = self.window.ui.documentTree.hasFocus() if focused: self.document.delete_components(self._selected_components()) def _delete_component(self, component: Component) -> None: self.document.delete_component(component) def _selected_components(self) -> list[Component]: indexes = self.window.ui.documentTree.selectionModel().selectedRows(0) selected = {id(component) for index in indexes if isinstance(component := self.model.value(index), Component)} components: list[Component] = [] for index in indexes: component = self.model.value(index) if not isinstance(component, Component): continue parent = index.parent() nested = False while parent.isValid(): value = self.model.value(parent) if isinstance(value, Component) and id(value) in selected: nested = True break parent = parent.parent() if not nested: components.append(component) return components