feat: add loader/helper for widget plugins

bec_lib/bec_lib/plugin_helper.py

@@ -51,6 +51,19 @@ def get_plugin_class(class_spec: str, additional_modules=None) -> type:
     return getattr(module, class_name)
 
 
+def get_widget_plugins() -> dict[str, type]:
+    """
+    Load all widget plugins.
+
+    Returns:
+        dict: A dictionary with the plugin names as keys and the plugin classes as values.
+    """
+    if not (widget_plugin := _get_available_plugins("bec.widgets.user_widgets")):
+        return {}
+    # Implemented in the plugin repo to avoid introducing a dependency on bec_widgets here.
+    return widget_plugin[0].get_all_plugin_widgets()
+
+
 def get_scan_plugins() -> dict:
     """
     Load all scan plugins.
@@ -88,7 +101,7 @@ def get_file_writer_plugins() -> dict:
 
 
 @cache
-def get_metadata_schema_registry() -> tuple[dict, type[BasicScanMetadata]]:
+def get_metadata_schema_registry() -> tuple[dict, type[BasicScanMetadata] | None]:
     module = _get_available_plugins("bec.scans.metadata_schema")
     if len(module) == 0:
         return {}, None
@@ -179,3 +192,7 @@ def _import_module(module_name: str):
     """
     module = importlib.import_module(module_name)
     return module
+
+
+if __name__ == "__main__":
+    plugs = get_widget_plugins()

bec_lib/bec_lib/scans.py

@@ -309,7 +309,7 @@ class Scans:
         params = {}
         for cmds in partition(bundle_size, args):
             cmds_serialized = [
-                cmd._compile_function_path() if hasattr(cmd, "_compile_function_path") else cmd
+                (cmd._compile_function_path() if hasattr(cmd, "_compile_function_path") else cmd)
                 for cmd in cmds
             ]
             params[cmds_serialized[0]] = cmds_serialized[1:]

bec_lib/util_scripts/create_plugin_structure.py

@@ -137,7 +137,8 @@ class PluginStructure:
 
         self.create_dir(f"{self.plugin_name}/bec_widgets/configs")
         self.create_dir(f"{self.plugin_name}/bec_widgets/widgets")
-        self.create_init_file(f"{self.plugin_name}/bec_widgets/widgets")
+        widget_init = os.path.join(current_dir, "plugin_setup_files", "widgets_init_file.py")
+        os.system(f"cp {widget_init} {self.plugin_name}/bec_widgets/widgets/__init__.py")
 
     def add_file_writer(self):
         self.create_dir(f"{self.plugin_name}/file_writer")

bec_lib/util_scripts/plugin_setup_files/widgets_init_file.py

@@ -0,0 +1,53 @@
+"""Module for user-created widget files and some utilities to load them."""
+
+import inspect
+import pkgutil
+from importlib import util as importlib_util
+from importlib.machinery import FileFinder, SourceFileLoader
+from types import ModuleType
+
+from bec_widgets import BECWidget
+from my_plugin_repo.bec_widgets import widgets
+
+
+def get_all_plugin_widgets() -> dict[str, BECWidget]:
+    """BEC uses this function to load widgets from this plugin repository. Don't modify it unless
+    you are absolutely certain of what you are doing!"""
+    return _all_widgets_from_all_submodules(widgets)
+
+
+def _get_widgets_from_module(module: ModuleType) -> dict[str, BECWidget]:
+    """Find any BECWidget subclasses in the given module and return them with their names."""
+    return dict(
+        inspect.getmembers(
+            module,
+            predicate=lambda item: inspect.isclass(item)
+            and issubclass(item, BECWidget)
+            and item is not BECWidget,
+        )
+    )
+
+
+def _all_widgets_from_all_submodules(module):
+    """Recursively load submodules, find any BECWidgets, and return them all as a flat dict."""
+    widgets = _get_widgets_from_module(module)
+    if not hasattr(module, "__path__"):
+        return widgets
+    submodule_specs = (
+        module_info.module_finder.find_spec(module_info.name)
+        for module_info in pkgutil.iter_modules(module.__path__)
+        if isinstance(module_info.module_finder, FileFinder)
+    )
+    for submodule in (
+        importlib_util.module_from_spec(spec) for spec in submodule_specs if spec is not None
+    ):
+        assert isinstance(
+            submodule.__loader__, SourceFileLoader
+        ), "Module found from FileFinder should have SourceFileLoader!"
+        submodule.__loader__.exec_module(submodule)
+        widgets.update(_all_widgets_from_all_submodules(submodule))
+    return widgets
+
+
+if __name__ == "__main__":  # pragma: no cover
+    print(get_all_plugin_widgets())