from collections.abc import Callable from functools import partial from typing import Protocol from PySide6.QtCore import QEvent, QItemSelectionModel, QObject, QPoint, QSize, Qt from PySide6.QtGui import QMouseEvent from PySide6.QtWidgets import QAbstractItemView, QDialog, QHeaderView, QMenu from bedit_core.models import Component, ComponentID, ConnectionID, 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 Graph, GraphComponentLabel, Icon, PortMetadata 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_fitted_icon ICON_SIZE = QSize(16, 16) class InterfaceEditorLike(Protocol): def exec(self) -> int: ... def ports(self) -> dict[PortID, Port]: ... def port_metadata(self) -> dict[PortID, PortMetadata]: ... class ParamEditorLike(Protocol): def exec(self) -> int: ... def params(self) -> dict[ParameterID, Parameter]: ... InterfaceEditorFactory = Callable[ [dict[PortID, Port], dict[PortID, PortMetadata], 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) window.equation_editor.equation_text_change_requested.connect(self.document.update_component_equation_text) window.equation_editor.port_metadata_change_requested.connect(self._change_equation_port_metadata) document.model_changed.connect(self._on_document_changed) document.icon_changed.connect(self._on_icon_changed) document.port_metadata_database_changed.connect(self._on_port_metadata_database_changed) document.graph_component_position_changed.connect(self._on_graph_component_position_changed) document.graph_component_label_changed.connect(self._on_graph_component_label_changed) document.graph_connection_points_changed.connect(self._on_graph_connection_points_changed) document.equation_text_changed.connect(self._on_equation_text_changed) self.model.rename_document_requested.connect(self.document.rename) self.model.rename_component_requested.connect(self.document.rename_component) window.graph_editor.component_moves_requested.connect(self.document.move_graph_components) window.graph_editor.component_label_move_requested.connect(self.document.move_graph_component_label) window.graph_editor.component_context_menu_requested.connect(self._show_graph_component_context_menu) window.graph_editor.component_open_requested.connect(self._open_graph_component) window.graph_editor.connection_points_change_requested.connect(self.document.change_graph_connection_points) window.graph_editor.connection_add_requested.connect(self.document.add_graph_connection) window.graph_editor.connections_delete_requested.connect(self.document.delete_graph_connections) # 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(24, 24)) 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, 28) self._tree_viewport = window.ui.documentTree.viewport() self._tree_viewport.installEventFilter(self) self._on_document_changed(document.model) def eventFilter(self, watched: QObject, event: QEvent) -> bool: if watched is self._tree_viewport and event.type() in (QEvent.Type.MouseButtonPress, QEvent.Type.MouseButtonRelease): assert isinstance(event, QMouseEvent) if event.button() == Qt.MouseButton.RightButton: if event.type() == QEvent.Type.MouseButtonRelease: self._show_context_menu(event.position().toPoint()) return True return super().eventFilter(watched, event) def _on_document_changed(self, model: CoreDocument) -> None: """Rebuild the tree whenever New/Open replaces the core document.""" displayed_component = self.window.graph_editor.component() or self.window.equation_editor.component() 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_fitted_icon(icon, component.interface.ports, ICON_SIZE)) self._show_component(displayed_component if any(component is displayed_component for component in self._components.values()) else None) # 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_component(component if isinstance(component, Component) else None) def _show_component(self, component: Component | None) -> None: if component is not None and isinstance(component.implementation, EquationImplementation): port_metadata = {port_id: self.document.port_metadata(port_id) for port_id in component.interface.ports} self.window.equation_editor.set_component(component, port_metadata) self.window.equation_editor.show() else: self.window.equation_editor.set_component(None) self.window.equation_editor.hide() if component is not None and isinstance(component.implementation, GraphImplementation): graph = self.document.graph_database().graphs.get(self.document.component_id(component), Graph()) icons = {component_id: self.document.component_icon(component_id) for component_id in component.implementation.graph.components} port_metadata = {port_id: self.document.port_metadata(port_id) for child in component.implementation.graph.components.values() for port_id in child.interface.ports} self.window.graph_editor.set_component(component, graph, icons, port_metadata) self.window.graph_editor.show() else: self.window.graph_editor.set_component(None) self.window.graph_editor.hide() def _on_equation_text_changed(self, component: Component, section: str) -> None: if self.window.equation_editor.component() is component: self.window.equation_editor.refresh_text(section) 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_fitted_icon(icon if isinstance(icon, Icon) else Icon(), component.interface.ports, ICON_SIZE)) graph_component = self.window.graph_editor.component() if graph_component is not None and isinstance(graph_component.implementation, GraphImplementation) and component_id in graph_component.implementation.graph.components: self._show_component(graph_component) def _on_port_metadata_database_changed(self, _database: object) -> None: graph_component = self.window.graph_editor.component() if graph_component is not None: self._show_component(graph_component) equation_component = self.window.equation_editor.component() if equation_component is not None: self.window.equation_editor.refresh_port_metadata({port_id: self.document.port_metadata(port_id) for port_id in equation_component.interface.ports}) def _change_equation_port_metadata(self, component: Component, port_metadata: dict[PortID, PortMetadata]) -> None: self.document.update_component_ports(component, component.interface.ports, port_metadata) def _on_graph_component_position_changed(self, graph_id: ComponentID, component_id: ComponentID, position: tuple[int, int] | None) -> None: graph_component = self.window.graph_editor.component() if graph_component is not None and self.document.component_id(graph_component) == graph_id: self.window.graph_editor.set_component_position(component_id, position) def _on_graph_component_label_changed(self, graph_id: ComponentID, component_id: ComponentID, label: object) -> None: graph_component = self.window.graph_editor.component() if graph_component is not None and self.document.component_id(graph_component) == graph_id: self.window.graph_editor.set_component_label(component_id, label if isinstance(label, GraphComponentLabel) else None) def _on_graph_connection_points_changed(self, graph_id: ComponentID, connection_id: ConnectionID, points: list[tuple[int, int]] | None) -> None: graph_component = self.window.graph_editor.component() if graph_component is not None and self.document.component_id(graph_component) == graph_id: self.window.graph_editor.set_connection_points(connection_id, points) 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) value = self.model.value(index) global_position = self.window.ui.documentTree.viewport().mapToGlobal(position) if isinstance(value, CoreDocument): self._show_root_context_menu(global_position) elif isinstance(value, Component): self._show_component_context_menu(value, global_position) def _show_root_context_menu(self, global_position: QPoint) -> None: menu = QMenu(self.window.ui.documentTree) add_graph_component = menu.addAction("Add Graph Component") add_equation_component = menu.addAction("Add Equation Component") selected = menu.exec(global_position) if selected is add_graph_component: self.document.add_empty_root_graph_component() elif selected is add_equation_component: self.document.add_empty_root_equation_component() def _show_graph_component_context_menu(self, component_id: ComponentID, global_position: QPoint) -> None: component = self._components.get(component_id) if component is not None: self._show_component_context_menu(component, global_position, component_id) def _open_graph_component(self, component_id: ComponentID) -> None: index = self.model.component_index(component_id) if index.isValid(): self.window.ui.documentTree.selectionModel().setCurrentIndex(index, QItemSelectionModel.SelectionFlag.ClearAndSelect | QItemSelectionModel.SelectionFlag.Rows) self.window.ui.documentTree.scrollTo(index) def _show_component_context_menu(self, component: Component, global_position: QPoint, graph_component_id: ComponentID | None = None) -> None: menu = QMenu(self.window.ui.documentTree) edit_interface = menu.addAction("Edit Interface") edit_params = menu.addAction("Edit Parameters") edit_icon = menu.addAction("Edit Icon") show_label = None if graph_component_id is not None: show_label = menu.addAction("Show Label") show_label.setCheckable(True) show_label.setChecked(self.window.graph_editor.component_label_visible(graph_component_id)) menu.addSeparator() add_graph_component = None add_equation_component = None if isinstance(component.implementation, GraphImplementation): 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(global_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 show_label is not None and selected is show_label: graph_component = self.window.graph_editor.component() if graph_component is not None: self.document.set_graph_component_label_visible(graph_component, graph_component_id, show_label.isChecked()) elif add_graph_component is not None and selected is add_graph_component: self._add_graph_component(component) elif add_equation_component is not None and 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, {port_id: self.document.port_metadata(port_id) for port_id in component.interface.ports}, self.window, ) if dialog.exec() == QDialog.DialogCode.Accepted: self.document.update_component_ports(component, dialog.ports(), dialog.port_metadata()) 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