From ef25f5638032f931ceb292540ada618508bb2aed Mon Sep 17 00:00:00 2001
From: wakonig_k <klaus.wakonig@psi.ch>
Date: Wed, 19 Jun 2024 15:36:05 +0200
Subject: [PATCH] feat(rpc): discover widgets automatically

---
 bec_widgets/cli/rpc_wigdet_handler.py       | 46 +++++++++++++++------
 bec_widgets/widgets/dock/dock.py            |  6 +--
 tests/unit_tests/test_rpc_widget_handler.py |  7 ++++
 3 files changed, 43 insertions(+), 16 deletions(-)
 create mode 100644 tests/unit_tests/test_rpc_widget_handler.py

diff --git a/bec_widgets/cli/rpc_wigdet_handler.py b/bec_widgets/cli/rpc_wigdet_handler.py
index ff2f372e..44c1d61f 100644
--- a/bec_widgets/cli/rpc_wigdet_handler.py
+++ b/bec_widgets/cli/rpc_wigdet_handler.py
@@ -1,22 +1,37 @@
 from bec_widgets.utils import BECConnector
-from bec_widgets.widgets.figure import BECFigure
-from bec_widgets.widgets.spiral_progress_bar.spiral_progress_bar import SpiralProgressBar
-from bec_widgets.widgets.text_box.text_box import TextBox
-from bec_widgets.widgets.website.website import WebsiteWidget
 
 
 class RPCWidgetHandler:
     """Handler class for creating widgets from RPC messages."""
 
-    widget_classes = {
-        "BECFigure": BECFigure,
-        "SpiralProgressBar": SpiralProgressBar,
-        "Website": WebsiteWidget,
-        "TextBox": TextBox,
-    }
+    def __init__(self):
+        self._widget_classes = None
 
-    @staticmethod
-    def create_widget(widget_type, **kwargs) -> BECConnector:
+    @property
+    def widget_classes(self):
+        """
+        Get the available widget classes.
+
+        Returns:
+            dict: The available widget classes.
+        """
+        if self._widget_classes is None:
+            self.update_available_widgets()
+        return self._widget_classes
+
+    def update_available_widgets(self):
+        """
+        Update the available widgets.
+
+        Returns:
+            None
+        """
+        from bec_widgets.utils.plugin_utils import get_rpc_classes
+
+        clss = get_rpc_classes("bec_widgets")
+        self._widget_classes = {cls.__name__: cls for cls in clss["top_level_classes"]}
+
+    def create_widget(self, widget_type, **kwargs) -> BECConnector:
         """
         Create a widget from an RPC message.
 
@@ -27,7 +42,12 @@ class RPCWidgetHandler:
         Returns:
             widget(BECConnector): The created widget.
         """
-        widget_class = RPCWidgetHandler.widget_classes.get(widget_type)
+        if self._widget_classes is None:
+            self.update_available_widgets()
+        widget_class = self._widget_classes.get(widget_type)
         if widget_class:
             return widget_class(**kwargs)
         raise ValueError(f"Unknown widget type: {widget_type}")
+
+
+widget_handler = RPCWidgetHandler()
diff --git a/bec_widgets/widgets/dock/dock.py b/bec_widgets/widgets/dock/dock.py
index 7f5ab8a3..38520f54 100644
--- a/bec_widgets/widgets/dock/dock.py
+++ b/bec_widgets/widgets/dock/dock.py
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional
 from pydantic import Field
 from pyqtgraph.dockarea import Dock
 
-from bec_widgets.cli.rpc_wigdet_handler import RPCWidgetHandler
+from bec_widgets.cli.rpc_wigdet_handler import widget_handler
 from bec_widgets.utils import BECConnector, ConnectionConfig, GridLayoutManager
 
 if TYPE_CHECKING:
@@ -149,7 +149,7 @@ class BECDock(BECConnector, Dock):
         Returns:
             list: The list of eligible widgets.
         """
-        return list(RPCWidgetHandler.widget_classes.keys())
+        return list(widget_handler.widget_classes.keys())
 
     def add_widget(
         self,
@@ -178,7 +178,7 @@ class BECDock(BECConnector, Dock):
             self.layout_manager.shift_widgets(shift, start_row=row)
 
         if isinstance(widget, str):
-            widget = RPCWidgetHandler.create_widget(widget)
+            widget = widget_handler.create_widget(widget)
         else:
             widget = widget
 
diff --git a/tests/unit_tests/test_rpc_widget_handler.py b/tests/unit_tests/test_rpc_widget_handler.py
new file mode 100644
index 00000000..f86d39e8
--- /dev/null
+++ b/tests/unit_tests/test_rpc_widget_handler.py
@@ -0,0 +1,7 @@
+from bec_widgets.cli.rpc_wigdet_handler import RPCWidgetHandler
+
+
+def test_rpc_widget_handler():
+    handler = RPCWidgetHandler()
+    assert "BECFigure" in handler.widget_classes
+    assert "SpiralProgressBar" in handler.widget_classes