bec-widgets 1.25.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

bec_widgets/cli/client_utils.py

@@ -1,38 +1,43 @@
+"""Client utilities for the BEC GUI."""
+
 from __future__ import annotations
 
-import importlib
-import importlib.metadata as imd
 import json
 import os
 import select
 import subprocess
 import threading
+import time
 from contextlib import contextmanager
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from threading import Lock
+from typing import TYPE_CHECKING, Literal, TypeAlias, cast
 
 from bec_lib.endpoints import MessageEndpoints
 from bec_lib.logger import bec_logger
-from bec_lib.utils.import_utils import isinstance_based_on_class_name, lazy_import, lazy_import_from
+from bec_lib.utils.import_utils import lazy_import_from
+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
-
-if TYPE_CHECKING:
-    from bec_lib import messages
-    from bec_lib.connector import MessageObject
-    from bec_lib.device import DeviceBase
+from bec_widgets.cli.rpc.rpc_base import RPCBase, RPCReference
+from bec_widgets.utils.serialization import register_serializer_extension
 
-    from bec_widgets.utils.bec_dispatcher import BECDispatcher
+if TYPE_CHECKING:  # pragma: no cover
+    from bec_lib.messages import GUIRegistryStateMessage
 else:
-    messages = lazy_import("bec_lib.messages")
-    # from bec_lib.connector import MessageObject
-    MessageObject = lazy_import_from("bec_lib.connector", ("MessageObject",))
-    BECDispatcher = lazy_import_from("bec_widgets.utils.bec_dispatcher", ("BECDispatcher",))
+    GUIRegistryStateMessage = lazy_import_from("bec_lib.messages", "GUIRegistryStateMessage")
 
 logger = bec_logger.logger
 
+IGNORE_WIDGETS = ["LaunchWindow"]
+
+RegistryState: TypeAlias = dict[
+    Literal["gui_id", "name", "widget_class", "config", "__rpc__", "container_proxy"],
+    str | bool | dict,
+]
+
+# pylint: disable=redefined-outer-scope
+
 
 def _filter_output(output: str) -> str:
     """
@@ -67,7 +72,13 @@ def _get_output(process, logger) -> None:
         logger.error(f"Error reading process output: {str(e)}")
 
 
-def _start_plot_process(gui_id: str, gui_class: type, config: dict | str, logger=None) -> None:
+def _start_plot_process(
+    gui_id: str,
+    gui_class_id: str,
+    config: dict | str,
+    gui_class: str = "dock_area",
+    logger=None,  # FIXME change gui_class back to "launcher" later
+) -> tuple[subprocess.Popen[str], threading.Thread | None]:
     """
     Start the plot in a new process.
 
@@ -76,7 +87,16 @@ def _start_plot_process(gui_id: str, gui_class: type, config: dict | str, logger
     process will not be captured.
     """
     # pylint: disable=subprocess-run-check
-    command = ["bec-gui-server", "--id", gui_id, "--gui_class", gui_class.__name__, "--hide"]
+    command = [
+        "bec-gui-server",
+        "--id",
+        gui_id,
+        "--gui_class",
+        gui_class,
+        "--gui_class_id",
+        gui_class_id,
+        "--hide",
+    ]
     if config:
         if isinstance(config, dict):
             config = json.dumps(config)
@@ -111,16 +131,20 @@ def _start_plot_process(gui_id: str, gui_class: type, config: dict | str, logger
 
 
 class RepeatTimer(threading.Timer):
+    """RepeatTimer class."""
+
     def run(self):
         while not self.finished.wait(self.interval):
             self.function(*self.args, **self.kwargs)
 
 
+# pylint: disable=protected-access
 @contextmanager
-def wait_for_server(client):
+def wait_for_server(client: BECGuiClient):
+    """Context manager to wait for the server to start."""
     timeout = client._startup_timeout
     if not timeout:
-        if client.gui_is_alive():
+        if client._gui_is_alive():
             # there is hope, let's wait a bit
             timeout = 1
         else:
@@ -138,132 +162,223 @@ def wait_for_server(client):
     yield
 
 
-### ----------------------------
-### NOTE
-### it is far easier to extend the 'delete' method on the client side,
-### to know when the client is deleted, rather than listening to server
-### to get notified. However, 'generate_cli.py' cannot add extra stuff
-### in the generated client module. So, here a class with the same name
-### is created, and client module is patched.
-class BECDockArea(client.BECDockArea):
-    def delete(self):
-        if self is BECGuiClient._top_level["main"].widget:
-            raise RuntimeError("Cannot delete main window")
-        super().delete()
-        try:
-            del BECGuiClient._top_level[self._gui_id]
-        except KeyError:
-            # if a dock area is not at top level
-            pass
-
-
-client.BECDockArea = BECDockArea
-### ----------------------------
+class WidgetNameSpace:
+    def __repr__(self):
+        console = Console()
+        table = Table(title="Available widgets for BEC CLI usage")
+        table.add_column("Widget Name", justify="left", style="magenta")
+        table.add_column("Description", justify="left")
+        for attr, value in self.__dict__.items():
+            docs = value.__doc__
+            docs = docs if docs else "No description available"
+            table.add_row(attr, docs)
+        console.print(table)
+        return ""
 
 
-@dataclass
-class WidgetDesc:
-    title: str
-    widget: BECDockArea
+class AvailableWidgetsNamespace:
+    """Namespace for available widgets in the BEC GUI."""
+
+    def __init__(self):
+        for widget in client.Widgets:
+            name = widget.value
+            if name in IGNORE_WIDGETS:
+                continue
+            setattr(self, name, name)
+
+    def __repr__(self):
+        console = Console()
+        table = Table(title="Available widgets for BEC CLI usage")
+        table.add_column("Widget Name", justify="left", style="magenta")
+        table.add_column("Description", justify="left")
+        for attr_name, _ in self.__dict__.items():
+            docs = getattr(client, attr_name).__doc__
+            docs = docs if docs else "No description available"
+            table.add_row(attr_name, docs if len(docs.strip()) > 0 else "No description available")
+        console.print(table)
+        return ""
 
 
 class BECGuiClient(RPCBase):
-    _top_level = {}
+    """BEC GUI client class. Container for GUI applications within Python."""
 
     def __init__(self, **kwargs) -> None:
         super().__init__(**kwargs)
-        self._auto_updates_enabled = True
-        self._auto_updates = None
+        self._lock = Lock()
+        self._anchor_widget = "launcher"
+        self._killed = False
+        self._top_level: dict[str, RPCReference] = {}
         self._startup_timeout = 0
         self._gui_started_timer = None
         self._gui_started_event = threading.Event()
         self._process = None
         self._process_output_processing_thread = None
+        self._server_registry: dict[str, RegistryState] = {}
+        self._ipython_registry: dict[str, RPCReference] = {}
+        self.available_widgets = AvailableWidgetsNamespace()
+        register_serializer_extension()
 
-    @property
-    def windows(self):
-        return self._top_level
+    ####################
+    #### Client API ####
+    ####################
 
     @property
-    def auto_updates(self):
-        if self._auto_updates_enabled:
-            with wait_for_server(self):
-                return self._auto_updates
+    def launcher(self) -> RPCBase:
+        """The launcher object."""
+        return RPCBase(gui_id=f"{self._gui_id}:launcher", parent=self, object_name="launcher")
+
+    def connect_to_gui_server(self, gui_id: str) -> None:
+        """Connect to a GUI server"""
+        # Unregister the old callback
+        self._client.connector.unregister(
+            MessageEndpoints.gui_registry_state(self._gui_id), cb=self._handle_registry_update
+        )
+        self._gui_id = gui_id
+
+        # reset the namespace
+        self._update_dynamic_namespace({})
+        self._server_registry = {}
+        self._top_level = {}
+        self._ipython_registry = {}
+
+        # Register the new callback
+        self._client.connector.register(
+            MessageEndpoints.gui_registry_state(self._gui_id),
+            cb=self._handle_registry_update,
+            parent=self,
+            from_start=True,
+        )
 
-    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"].widget)
-                except Exception as e:
-                    logger.error(f"Error loading auto update script from plugin: {str(e)}")
-        return None
+    @property
+    def windows(self) -> dict:
+        """Dictionary with dock areas in the GUI."""
+        return {widget.object_name: widget for widget in self._top_level.values()}
 
     @property
-    def selected_device(self):
+    def window_list(self) -> list:
+        """List with dock areas in the GUI."""
+        return list(self._top_level.values())
+
+    def start(self, wait: bool = False) -> None:
+        """Start the GUI server."""
+        return self._start(wait=wait)
+
+    def show(self):
+        """Show the GUI window."""
+        if self._check_if_server_is_alive():
+            return self._show_all()
+        return self.start(wait=True)
+
+    def hide(self):
+        """Hide the GUI window."""
+        return self._hide_all()
+
+    def new(
+        self,
+        name: str | None = None,
+        wait: bool = True,
+        geometry: tuple[int, int, int, int] | None = None,
+        launch_script: str = "dock_area",
+        **kwargs,
+    ) -> client.BECDockArea:
+        """Create a new top-level dock area.
+
+        Args:
+            name(str, optional): The name of the dock area. Defaults to None.
+            wait(bool, optional): Whether to wait for the server to start. Defaults to True.
+            geometry(tuple[int, int, int, int] | None): The geometry of the dock area (pos_x, pos_y, w, h)
+        Returns:
+            client.BECDockArea: The new dock area.
         """
-        Selected device for the plot.
+        if not self._check_if_server_is_alive():
+            self.start(wait=True)
+        if wait:
+            with wait_for_server(self):
+                widget = self.launcher._run_rpc(
+                    "launch", launch_script=launch_script, name=name, geometry=geometry, **kwargs
+                )  # pylint: disable=protected-access
+                return widget
+        widget = self.launcher._run_rpc(
+            "launch", launch_script=launch_script, name=name, geometry=geometry, **kwargs
+        )  # pylint: disable=protected-access
+        return widget
+
+    def delete(self, name: str) -> None:
+        """Delete a dock area.
+
+        Args:
+            name(str): The name of the dock area.
         """
-        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")
+        widget = self.windows.get(name)
+        if widget is None:
+            raise ValueError(f"Dock area {name} not found.")
+        widget._run_rpc("close")  # pylint: disable=protected-access
 
-    def _start_update_script(self) -> None:
-        self._client.connector.register(MessageEndpoints.scan_status(), cb=self._handle_msg_update)
+    def delete_all(self) -> None:
+        """Delete all dock areas."""
+        for widget_name in self.windows:
+            self.delete(widget_name)
 
-    def _handle_msg_update(self, msg: MessageObject) -> None:
-        if self.auto_updates is not None:
-            # pylint: disable=protected-access
-            return self._update_script_msg_parser(msg.value)
+    def kill_server(self) -> None:
+        """Kill the GUI server."""
+        # Unregister the registry state
+        self._killed = True
 
-    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 self._gui_started_timer is not None:
+            self._gui_started_timer.cancel()
+            self._gui_started_timer.join()
 
-    def _gui_post_startup(self):
-        self._top_level["main"] = WidgetDesc(
-            title="BEC Widgets", widget=BECDockArea(gui_id=self._gui_id)
+        if self._process is None:
+            return
+
+        if self._process:
+            logger.success("Stopping GUI...")
+            self._process.terminate()
+            if self._process_output_processing_thread:
+                self._process_output_processing_thread.join()
+            self._process.wait()
+            self._process = None
+
+        # Unregister the registry state
+        self._client.connector.unregister(
+            MessageEndpoints.gui_registry_state(self._gui_id), cb=self._handle_registry_update
         )
-        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()
+        # Remove all reference from top level
+        self._top_level.clear()
+        self._server_registry.clear()
+
+    def close(self):
+        """Deprecated. Use kill_server() instead."""
+        # FIXME, deprecated in favor of kill, will be removed in the future
+        self.kill_server()
+
+    #########################
+    #### Private methods ####
+    #########################
+
+    def _check_if_server_is_alive(self):
+        """Checks if the process is alive"""
+        if self._process is None:
+            return False
+        if self._process.poll() is not None:
+            return False
+        return True
+
+    def _gui_post_startup(self):
+        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._server_registry.keys())) < 2 or not hasattr(
+                self, self._anchor_widget
+            ):
+                time.sleep(0.1)
+            else:
+                break
+
         self._gui_started_event.set()
 
-    def start_server(self, wait=False) -> None:
+    def _start_server(self, wait: bool = False) -> None:
         """
         Start the GUI server, and execute callback when it is launched
         """
@@ -272,7 +387,10 @@ class BECGuiClient(RPCBase):
             self._startup_timeout = 5
             self._gui_started_event.clear()
             self._process, self._process_output_processing_thread = _start_plot_process(
-                self._gui_id, self.__class__, self._client._service_config.config, logger=logger
+                self._gui_id,
+                gui_class_id="bec",
+                config=self._client._service_config.config,  # pylint: disable=protected-access
+                logger=logger,
             )
 
             def gui_started_callback(callback):
@@ -280,80 +398,134 @@ class BECGuiClient(RPCBase):
                     if callable(callback):
                         callback()
                 finally:
-                    threading.current_thread().cancel()
+                    threading.current_thread().cancel()  # type: ignore
 
             self._gui_started_timer = RepeatTimer(
-                0.5, lambda: self.gui_is_alive() and gui_started_callback(self._gui_post_startup)
+                0.5, lambda: self._gui_is_alive() and gui_started_callback(self._gui_post_startup)
             )
             self._gui_started_timer.start()
 
         if wait:
             self._gui_started_event.wait()
 
-    def _dump(self):
-        rpc_client = RPCBase(gui_id=f"{self._gui_id}:window", parent=self)
-        return rpc_client._run_rpc("_dump")
-
-    def start(self):
-        return self.start_server()
+    def _start(self, wait: bool = False) -> None:
+        self._killed = False
+        self._client.connector.register(
+            MessageEndpoints.gui_registry_state(self._gui_id),
+            cb=self._handle_registry_update,
+            parent=self,
+        )
+        return self._start_server(wait=wait)
+
+    @staticmethod
+    def _handle_registry_update(
+        msg: dict[str, GUIRegistryStateMessage], parent: BECGuiClient
+    ) -> None:
+        # This was causing a deadlock during shutdown, not sure why.
+        # with self._lock:
+        self = parent
+        self._server_registry = cast(dict[str, RegistryState], msg["data"].state)
+        self._update_dynamic_namespace(self._server_registry)
 
     def _do_show_all(self):
-        rpc_client = RPCBase(gui_id=f"{self._gui_id}:window", parent=self)
-        rpc_client._run_rpc("show")
+        rpc_client = RPCBase(gui_id=f"{self._gui_id}:launcher", parent=self)
+        rpc_client._run_rpc("show")  # pylint: disable=protected-access
         for window in self._top_level.values():
-            window.widget.show()
+            window.show()
 
-    def show_all(self):
+    def _show_all(self):
         with wait_for_server(self):
             return self._do_show_all()
 
-    def hide_all(self):
-        with wait_for_server(self):
-            rpc_client = RPCBase(gui_id=f"{self._gui_id}:window", parent=self)
-            rpc_client._run_rpc("hide")
-            for window in self._top_level.values():
-                window.widget.hide()
-
-    def show(self):
-        if self._process is not None:
-            return self.show_all()
-        # backward compatibility: show() was also starting server
-        return self.start_server(wait=True)
-
-    def hide(self):
-        return self.hide_all()
-
-    @property
-    def main(self):
-        """Return client to main dock area (in main window)"""
+    def _hide_all(self):
         with wait_for_server(self):
-            return self._top_level["main"].widget
+            rpc_client = RPCBase(gui_id=f"{self._gui_id}:launcher", parent=self)
+            rpc_client._run_rpc("hide")  # pylint: disable=protected-access
+            if not self._killed:
+                for window in self._top_level.values():
+                    window.hide()
 
-    def new(self, title):
-        """Ask main window to create a new top-level dock area"""
-        with wait_for_server(self):
-            rpc_client = RPCBase(gui_id=f"{self._gui_id}:window", parent=self)
-            widget = rpc_client._run_rpc("new_dock_area", title)
-            self._top_level[widget._gui_id] = WidgetDesc(title=title, widget=widget)
-            return widget
+    def _update_dynamic_namespace(self, server_registry: dict):
+        """
+        Update the dynamic name space with the given server registry.
+        Setting the server registry to an empty dictionary will remove all widgets from the namespace.
 
-    def close(self) -> None:
+        Args:
+            server_registry (dict): The server registry
         """
-        Close the gui window.
+        top_level_widgets: dict[str, RPCReference] = {}
+        for gui_id, state in server_registry.items():
+            widget = self._add_widget(state, self)
+            if widget is None:
+                # ignore widgets that are not supported
+                continue
+            # get all top-level widgets. These are widgets that have no parent
+            if not state["config"].get("parent_id"):
+                top_level_widgets[gui_id] = widget
+
+        remove_from_registry = []
+        for gui_id, widget in self._ipython_registry.items():
+            if gui_id not in server_registry:
+                remove_from_registry.append(gui_id)
+        for gui_id in remove_from_registry:
+            self._ipython_registry.pop(gui_id)
+
+        removed_widgets = [
+            widget.object_name for widget in self._top_level.values() if widget._is_deleted()
+        ]
+
+        for widget_name in removed_widgets:
+            # the check is not strictly necessary, but better safe
+            # than sorry; who knows what the user has done
+            if hasattr(self, widget_name):
+                delattr(self, widget_name)
+
+        for gui_id, widget_ref in top_level_widgets.items():
+            setattr(self, widget_ref.object_name, widget_ref)
+
+        self._top_level = top_level_widgets
+
+        for widget in self._ipython_registry.values():
+            widget._refresh_references()
+
+    def _add_widget(self, state: dict, parent: object) -> RPCReference | None:
+        """Add a widget to the namespace
+
+        Args:
+            state (dict): The state of the widget from the _server_registry.
+            parent (object): The parent object.
         """
-        self._top_level.clear()
+        object_name = state["object_name"]
+        gui_id = state["gui_id"]
+        if state["widget_class"] in IGNORE_WIDGETS:
+            return
+        widget_class = getattr(client, state["widget_class"], None)
+        if widget_class is None:
+            return
+        obj = self._ipython_registry.get(gui_id)
+        if obj is None:
+            widget = widget_class(gui_id=gui_id, object_name=object_name, parent=parent)
+            self._ipython_registry[gui_id] = widget
+        else:
+            widget = obj
+        obj = RPCReference(registry=self._ipython_registry, gui_id=gui_id)
+        return obj
 
-        if self._gui_started_timer is not None:
-            self._gui_started_timer.cancel()
-            self._gui_started_timer.join()
 
-        if self._process is None:
-            return
+if __name__ == "__main__":  # pragma: no cover
+    from bec_lib.client import BECClient
+    from bec_lib.service_config import ServiceConfig
 
-        if self._process:
-            logger.success("Stopping GUI...")
-            self._process.terminate()
-            if self._process_output_processing_thread:
-                self._process_output_processing_thread.join()
-            self._process.wait()
-            self._process = None
+    try:
+        config = ServiceConfig()
+        bec_client = BECClient(config)
+        bec_client.start()
+
+        # Test the client_utils.py module
+        gui = BECGuiClient()
+
+        gui.start(wait=True)
+        gui.new().new(widget="Waveform")
+        time.sleep(10)
+    finally:
+        gui.kill_server()