diff --git a/bec_lib/bec_lib/plugin_helper.py b/bec_lib/bec_lib/plugin_helper.py index 2308c12c..b9afd027 100644 --- a/bec_lib/bec_lib/plugin_helper.py +++ b/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() diff --git a/bec_lib/bec_lib/scans.py b/bec_lib/bec_lib/scans.py index 300c97ba..29dad7e9 100644 --- a/bec_lib/bec_lib/scans.py +++ b/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:] diff --git a/bec_lib/util_scripts/create_plugin_structure.py b/bec_lib/util_scripts/create_plugin_structure.py index 230d98ea..7c945454 100644 --- a/bec_lib/util_scripts/create_plugin_structure.py +++ b/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") diff --git a/bec_lib/util_scripts/plugin_setup_files/widgets_init_file.py b/bec_lib/util_scripts/plugin_setup_files/widgets_init_file.py new file mode 100644 index 00000000..89ff06e6 --- /dev/null +++ b/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())