From 7ba93ce934cc644ad6340f141a6a0888bd1d3d98 Mon Sep 17 00:00:00 2001 From: appel_c Date: Sat, 22 Mar 2025 07:36:13 +0100 Subject: [PATCH] refactor: cleanup rpc reference tracking, fix appquit, fix namespace updates edge cases --- bec_widgets/cli/client_utils.py | 125 +++--------------- bec_widgets/cli/rpc/rpc_base.py | 15 ++- bec_widgets/cli/rpc/rpc_register.py | 3 +- bec_widgets/cli/server.py | 7 + bec_widgets/utils/bec_connector.py | 1 - bec_widgets/utils/bec_widget.py | 7 +- bec_widgets/widgets/containers/dock/dock.py | 16 --- .../widgets/containers/dock/dock_area.py | 15 +-- 8 files changed, 51 insertions(+), 138 deletions(-) diff --git a/bec_widgets/cli/client_utils.py b/bec_widgets/cli/client_utils.py index a96a9efd..eec25511 100644 --- a/bec_widgets/cli/client_utils.py +++ b/bec_widgets/cli/client_utils.py @@ -2,8 +2,6 @@ from __future__ import annotations -import importlib -import importlib.metadata as imd import json import os import select @@ -21,24 +19,19 @@ from rich.console import Console from rich.table import Table import bec_widgets.cli.client as client - -# from bec_widgets.cli.auto_updates import AutoUpdates from bec_widgets.cli.rpc.rpc_base import RPCBase, RPCReference -if TYPE_CHECKING: - from bec_lib import messages - from bec_lib.connector import MessageObject - from bec_lib.device import DeviceBase +if TYPE_CHECKING: # pragma: no cover from bec_lib.redis_connector import StreamMessage else: - messages = lazy_import("bec_lib.messages") - MessageObject = lazy_import_from("bec_lib.connector", ("MessageObject",)) StreamMessage = lazy_import_from("bec_lib.redis_connector", ("StreamMessage",)) logger = bec_logger.logger IGNORE_WIDGETS = ["BECDockArea", "BECDock"] +# pylint: disable=redefined-outer-scope + def _filter_output(output: str) -> str: """ @@ -258,7 +251,7 @@ class BECGuiClient(RPCBase): """Show the GUI window.""" if self._check_if_server_is_alive(): return self._show_all() - return self._start(wait=True) + return self.start(wait=True) def hide(self): """Hide the GUI window.""" @@ -279,7 +272,8 @@ class BECGuiClient(RPCBase): Returns: client.BECDockArea: The new dock area. """ - self.show() + if not self._check_if_server_is_alive(): + self.start(wait=True) if wait: with wait_for_server(self): rpc_client = RPCBase(gui_id=f"{self._gui_id}:window", parent=self) @@ -313,11 +307,7 @@ class BECGuiClient(RPCBase): def kill_server(self) -> None: """Kill the GUI server.""" - self._top_level.clear() # Unregister the registry state - self._client.connector.unregister( - MessageEndpoints.gui_registry_state(self._gui_id), cb=self._handle_registry_update - ) self._killed = True if self._gui_started_timer is not None: @@ -339,6 +329,9 @@ class BECGuiClient(RPCBase): self._client.connector.unregister( MessageEndpoints.gui_registry_state(self._gui_id), cb=self._handle_registry_update ) + # Remove all reference from top level + self._top_level.clear() + self._registry_state.clear() def close(self): """Deprecated. Use kill_server() instead.""" @@ -358,24 +351,14 @@ class BECGuiClient(RPCBase): return True def _gui_post_startup(self): - timeout = 10 + timeout = 60 + # Wait for 'bec' gui to be registered, this may take some time + # After 60s timeout. Should this raise an exception on timeout? while time.time() < time.time() + timeout: if len(list(self._registry_state.keys())) == 0: time.sleep(0.1) else: break - # FIXME AUTO UPDATES - # if self._auto_updates_enabled: - # if self._auto_updates is None: - # auto_updates = self._get_update_script() - # if auto_updates is None: - # AutoUpdates.create_default_dock = True - # AutoUpdates.enabled = True - # auto_updates = AutoUpdates(self._top_level["main"].widget) - # if auto_updates.create_default_dock: - # auto_updates.start_default_dock() - # self._start_update_script() - # self._auto_updates = auto_updates self._do_show_all() self._gui_started_event.set() @@ -416,15 +399,19 @@ class BECGuiClient(RPCBase): def _start(self, wait: bool = False) -> None: self._killed = False + # Clear the registry state + self._registry_state.clear() + # Clear top level + self._top_level.clear() self._client.connector.register( MessageEndpoints.gui_registry_state(self._gui_id), cb=self._handle_registry_update ) return self._start_server(wait=wait) def _handle_registry_update(self, msg: StreamMessage) -> None: - # with self._lock: - self._registry_state = msg["data"].state - self._update_dynamic_namespace() + with self._lock: + self._registry_state = msg["data"].state + self._update_dynamic_namespace() def _do_show_all(self): rpc_client = RPCBase(gui_id=f"{self._gui_id}:window", parent=self) @@ -446,6 +433,8 @@ class BECGuiClient(RPCBase): def _update_dynamic_namespace(self): """Update the dynamic name space""" + # Clear the top level + self._top_level.clear() # First we update the name space based on the new registry state self._add_registry_to_namespace() # Then we clear the ipython registry from old objects @@ -559,75 +548,6 @@ class BECGuiClient(RPCBase): obj = RPCReference(registry=self._ipython_registry, gui_id=gui_id) return obj - ################################ - #### Auto updates #### - #### potentially deprecated #### - ################################ - - # FIXME AUTO UPDATES - # @property - # def auto_updates(self): - # if self._auto_updates_enabled: - # with wait_for_server(self): - # return self._auto_updates - - # def _get_update_script(self) -> AutoUpdates | None: - # eps = imd.entry_points(group="bec.widgets.auto_updates") - # for ep in eps: - # if ep.name == "plugin_widgets_update": - # try: - # spec = importlib.util.find_spec(ep.module) - # # if the module is not found, we skip it - # if spec is None: - # continue - # return ep.load()(gui=self._top_level["main"]) - # except Exception as e: - # logger.error(f"Error loading auto update script from plugin: {str(e)}") - # return None - - # FIXME AUTO UPDATES - # @property - # def selected_device(self) -> str | None: - # """ - # Selected device for the plot. - # """ - # auto_update_config_ep = MessageEndpoints.gui_auto_update_config(self._gui_id) - # auto_update_config = self._client.connector.get(auto_update_config_ep) - # if auto_update_config: - # return auto_update_config.selected_device - # return None - - # @selected_device.setter - # def selected_device(self, device: str | DeviceBase): - # if isinstance_based_on_class_name(device, "bec_lib.device.DeviceBase"): - # self._client.connector.set_and_publish( - # MessageEndpoints.gui_auto_update_config(self._gui_id), - # messages.GUIAutoUpdateConfigMessage(selected_device=device.name), - # ) - # elif isinstance(device, str): - # self._client.connector.set_and_publish( - # MessageEndpoints.gui_auto_update_config(self._gui_id), - # messages.GUIAutoUpdateConfigMessage(selected_device=device), - # ) - # else: - # raise ValueError("Device must be a string or a device object") - - # FIXME AUTO UPDATES - # def _start_update_script(self) -> None: - # self._client.connector.register(MessageEndpoints.scan_status(), cb=self._handle_msg_update) - - # def _handle_msg_update(self, msg: StreamMessage) -> None: - # if self.auto_updates is not None: - # # pylint: disable=protected-access - # return self._update_script_msg_parser(msg.value) - - # def _update_script_msg_parser(self, msg: messages.BECMessage) -> None: - # if isinstance(msg, messages.ScanStatusMessage): - # if not self._gui_is_alive(): - # return - # if self._auto_updates_enabled: - # return self.auto_updates.do_update(msg) - if __name__ == "__main__": # pragma: no cover from bec_lib.client import BECClient @@ -642,8 +562,7 @@ if __name__ == "__main__": # pragma: no cover gui = BECGuiClient() gui.start(wait=True) - print(gui.window_list) - gui.new() + gui.new().new(widget="Waveform") time.sleep(10) finally: gui.kill_server() diff --git a/bec_widgets/cli/rpc/rpc_base.py b/bec_widgets/cli/rpc/rpc_base.py index ab8fa51a..652337b1 100644 --- a/bec_widgets/cli/rpc/rpc_base.py +++ b/bec_widgets/cli/rpc/rpc_base.py @@ -4,7 +4,7 @@ import inspect import threading import uuid from functools import wraps -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any from bec_lib.client import BECClient from bec_lib.endpoints import MessageEndpoints @@ -20,6 +20,8 @@ else: # from bec_lib.connector import MessageObject MessageObject = lazy_import_from("bec_lib.connector", ("MessageObject",)) +# pylint: disable=protected-access + def rpc_call(func): """ @@ -40,6 +42,7 @@ def rpc_call(func): while caller_frame: if "jedi" in caller_frame.f_globals: # Jedi module is present, likely tab completion + # Do not run the RPC call return None # func(*args, **kwargs) caller_frame = caller_frame.f_back @@ -160,7 +163,7 @@ class RPCBase: parent = parent._parent return parent - def _run_rpc(self, method, *args, wait_for_rpc_response=True, timeout=3, **kwargs) -> Any: + def _run_rpc(self, method, *args, wait_for_rpc_response=True, timeout=300, **kwargs) -> Any: """ Run the RPC call. @@ -179,7 +182,6 @@ class RPCBase: parameter={"args": args, "kwargs": kwargs, "gui_id": self._gui_id}, metadata={"request_id": request_id}, ) - print(f"running and rpc {method}") # pylint: disable=protected-access receiver = self._root._gui_id if wait_for_rpc_response: @@ -233,7 +235,12 @@ class RPCBase: return msg_result cls = getattr(client, cls) - # print(msg_result) + # The namespace of the object will be updated dynamically on the client side + # Therefor it is important to check if the object is already in the registry + # If yes, we return the reference to the object, otherwise we create a new object + # pylint: disable=protected-access + if msg_result["gui_id"] in self._root._ipython_registry: + return RPCReference(self._root._ipython_registry, msg_result["gui_id"]) ret = cls(parent=self, **msg_result) self._root._ipython_registry[ret._gui_id] = ret obj = RPCReference(self._root._ipython_registry, ret._gui_id) diff --git a/bec_widgets/cli/rpc/rpc_register.py b/bec_widgets/cli/rpc/rpc_register.py index 719fa399..5f0d7a2f 100644 --- a/bec_widgets/cli/rpc/rpc_register.py +++ b/bec_widgets/cli/rpc/rpc_register.py @@ -37,7 +37,7 @@ def broadcast_update(func): @wraps(func) def wrapper(self, *args, **kwargs): result = func(self, *args, **kwargs) - self.broadcast() + # self.broadcast() return result return wrapper @@ -130,7 +130,6 @@ class RPCRegister: """ Broadcast the update to all the callbacks. """ - # print("Broadcasting") connections = self.list_all_connections() for callback in self.callbacks: callback(connections) diff --git a/bec_widgets/cli/server.py b/bec_widgets/cli/server.py index 101a3c57..c19ce60d 100644 --- a/bec_widgets/cli/server.py +++ b/bec_widgets/cli/server.py @@ -117,6 +117,7 @@ class BECWidgetsCLIServer: return obj def run_rpc(self, obj, method, args, kwargs): + # Run with rpc registry broadcast, but only once with rpc_register_broadcast(self.rpc_register): logger.debug(f"Running RPC instruction: {method} with args: {args}, kwargs: {kwargs}") method_obj = getattr(obj, method) @@ -315,6 +316,12 @@ def main(): def sigint_handler(*args): # display message, for people to let it terminate gracefully print("Caught SIGINT, exiting") + # Close all widgets + rpc_register = RPCRegister() + with rpc_register_broadcast(rpc_register): + for widget in QApplication.instance().topLevelWidgets(): + widget.close() + app.quit() signal.signal(signal.SIGINT, sigint_handler) diff --git a/bec_widgets/utils/bec_connector.py b/bec_widgets/utils/bec_connector.py index 443c2983..ec3082c1 100644 --- a/bec_widgets/utils/bec_connector.py +++ b/bec_widgets/utils/bec_connector.py @@ -318,7 +318,6 @@ class BECConnector: self.deleteLater() else: self.rpc_register.remove_rpc(self) - self.rpc_register.broadcast() def get_config(self, dict_output: bool = True) -> dict | BaseModel: """ diff --git a/bec_widgets/utils/bec_widget.py b/bec_widgets/utils/bec_widget.py index 2b649c4b..551e557d 100644 --- a/bec_widgets/utils/bec_widget.py +++ b/bec_widgets/utils/bec_widget.py @@ -7,9 +7,9 @@ from bec_lib.logger import bec_logger from qtpy.QtCore import Slot from qtpy.QtWidgets import QApplication, QWidget +from bec_widgets.cli.rpc.rpc_register import rpc_register_broadcast from bec_widgets.utils.bec_connector import BECConnector, ConnectionConfig from bec_widgets.utils.colors import set_theme -from bec_widgets.utils.container_utils import WidgetContainerUtils if TYPE_CHECKING: # pragma: no cover from bec_widgets.widgets.containers.dock import BECDock @@ -101,8 +101,9 @@ class BECWidget(BECConnector): def cleanup(self): """Cleanup the widget.""" - # All widgets need to call super().cleanup() in their cleanup method - self.rpc_register.remove_rpc(self) + with rpc_register_broadcast(self.rpc_register): + # All widgets need to call super().cleanup() in their cleanup method + self.rpc_register.remove_rpc(self) def closeEvent(self, event): """Wrap the close even to ensure the rpc_register is cleaned up.""" diff --git a/bec_widgets/widgets/containers/dock/dock.py b/bec_widgets/widgets/containers/dock/dock.py index 02a7618d..23251a30 100644 --- a/bec_widgets/widgets/containers/dock/dock.py +++ b/bec_widgets/widgets/containers/dock/dock.py @@ -8,7 +8,6 @@ from pyqtgraph.dockarea import Dock, DockLabel from qtpy import QtCore, QtGui from bec_widgets.cli.client_utils import IGNORE_WIDGETS -from bec_widgets.cli.rpc.rpc_register import RPCRegister from bec_widgets.cli.rpc.rpc_widget_handler import widget_handler from bec_widgets.utils import ConnectionConfig, GridLayoutManager from bec_widgets.utils.bec_widget import BECWidget @@ -336,13 +335,8 @@ class BECDock(BECWidget, Dock): if hasattr(widget, "config"): widget.config.gui_id = widget.gui_id self.config.widgets[widget._name] = widget.config # pylint: disable=protected-access - self._broadcast_update() return widget - def _broadcast_update(self): - rpc_register = RPCRegister() - rpc_register.broadcast() - def move_widget(self, widget: QWidget, new_row: int, new_col: int): """ Move a widget to a new position in the layout. @@ -400,7 +394,6 @@ class BECDock(BECWidget, Dock): if widget in self.widgets: self.widgets.remove(widget) widget.close() - # self._broadcast_update() def delete_all(self): """ @@ -423,15 +416,6 @@ class BECDock(BECWidget, Dock): self.label.deleteLater() super().cleanup() - # def closeEvent(self, event): # pylint: disable=uselsess-parent-delegation - # """Close Event for dock and cleanup. - - # This wrapper ensures that the BECWidget close event is triggered. - # If removed, the closeEvent from pyqtgraph will be triggered, which - # is not calling super().closeEvent(event) and will not trigger the BECWidget close event. - # """ - # return super().closeEvent(event) - def close(self): """ Close the dock area and cleanup. diff --git a/bec_widgets/widgets/containers/dock/dock_area.py b/bec_widgets/widgets/containers/dock/dock_area.py index 87a12a46..488378c7 100644 --- a/bec_widgets/widgets/containers/dock/dock_area.py +++ b/bec_widgets/widgets/containers/dock/dock_area.py @@ -11,7 +11,7 @@ from qtpy.QtCore import QSize, Qt from qtpy.QtGui import QPainter, QPaintEvent from qtpy.QtWidgets import QApplication, QSizePolicy, QVBoxLayout, QWidget -from bec_widgets.cli.rpc.rpc_register import RPCRegister +from bec_widgets.cli.rpc.rpc_register import RPCRegister, rpc_register_broadcast from bec_widgets.qt_utils.error_popups import SafeSlot from bec_widgets.qt_utils.toolbar import ( ExpandableMenuAction, @@ -225,8 +225,11 @@ class BECDockArea(BECWidget, QWidget): @SafeSlot() def _create_widget_from_toolbar(self, widget_name: str) -> None: - dock_name = WidgetContainerUtils.generate_unique_name(widget_name, self.panels.keys()) - self.new(name=dock_name, widget=widget_name) + rpc_register = RPCRegister() + # Run with RPC broadcast to namespace of all widgets + with rpc_register_broadcast(rpc_register): + dock_name = WidgetContainerUtils.generate_unique_name(widget_name, self.panels.keys()) + self.new(name=dock_name, widget=widget_name) def paintEvent(self, event: QPaintEvent): # TODO decide if we want any default instructions super().paintEvent(event) @@ -381,14 +384,8 @@ class BECDockArea(BECWidget, QWidget): self.update() if floating: dock.detach() - # Run broadcast update - self._broadcast_update() return dock - def _broadcast_update(self): - rpc_register = RPCRegister() - rpc_register.broadcast() - def detach_dock(self, dock_name: str) -> BECDock: """ Undock a dock from the dock area.