refactor: cleanup rpc reference tracking, fix appquit, fix namespace updates edge cases

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:
+if TYPE_CHECKING:  # pragma: no cover
-    from bec_lib import messages
-    from bec_lib.connector import MessageObject
-    from bec_lib.device import DeviceBase
     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:
+        with self._lock:
-        self._registry_state = msg["data"].state
+            self._registry_state = msg["data"].state
-        self._update_dynamic_namespace()
+            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().new(widget="Waveform")
-        gui.new()
         time.sleep(10)
     finally:
         gui.kill_server()

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)

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)

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)

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:
         """

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
+        with rpc_register_broadcast(self.rpc_register):
-        self.rpc_register.remove_rpc(self)
+            # 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."""

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.

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())
+        rpc_register = RPCRegister()
-        self.new(name=dock_name, widget=widget_name)
+        # 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.