From b225a7cc90b55697211c28d9411b6f85c8077217 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 27 May 2025 17:13:18 +0200 Subject: [PATCH] refactor: store modules with widget search --- bec_widgets/applications/launch_window.py | 2 +- bec_widgets/cli/client.py | 2 +- bec_widgets/cli/generate_cli.py | 2 +- bec_widgets/cli/rpc/rpc_widget_handler.py | 7 ++- bec_widgets/utils/bec_plugin_helper.py | 45 ++++++++++------ bec_widgets/utils/plugin_utils.py | 53 ++++++++++--------- bec_widgets/utils/ui_loader.py | 9 ++-- tests/unit_tests/test_bec_plugin_helper.py | 20 +++++-- .../unit_tests/test_client_plugin_widgets.py | 5 +- tests/unit_tests/test_generate_cli_client.py | 2 +- tests/unit_tests/test_rpc_widget_handler.py | 8 ++- tests/unit_tests/test_scan_control.py | 2 +- 12 files changed, 93 insertions(+), 64 deletions(-) diff --git a/bec_widgets/applications/launch_window.py b/bec_widgets/applications/launch_window.py index acada40d..4a62dc2f 100644 --- a/bec_widgets/applications/launch_window.py +++ b/bec_widgets/applications/launch_window.py @@ -242,7 +242,7 @@ class LaunchWindow(BECMainWindow): ) # plugin widgets - self.available_widgets: dict[str, BECWidget] = get_all_plugin_widgets() + self.available_widgets: dict[str, type[BECWidget]] = get_all_plugin_widgets().as_dict() if self.available_widgets: plugin_repo_name = next(iter(self.available_widgets.values())).__module__.split(".")[0] plugin_repo_name = plugin_repo_name.removesuffix("_bec").upper() diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py index c05d8d0d..22fb0b9d 100644 --- a/bec_widgets/cli/client.py +++ b/bec_widgets/cli/client.py @@ -63,7 +63,7 @@ _Widgets = { try: - _plugin_widgets = get_all_plugin_widgets() + _plugin_widgets = get_all_plugin_widgets().as_dict() plugin_client = get_plugin_client_module() Widgets = _WidgetsEnumType("Widgets", {name: name for name in _plugin_widgets} | _Widgets) diff --git a/bec_widgets/cli/generate_cli.py b/bec_widgets/cli/generate_cli.py index fab9a803..9c472c02 100644 --- a/bec_widgets/cli/generate_cli.py +++ b/bec_widgets/cli/generate_cli.py @@ -111,7 +111,7 @@ _Widgets = { self.content += """ try: - _plugin_widgets = get_all_plugin_widgets() + _plugin_widgets = get_all_plugin_widgets().as_dict() plugin_client = get_plugin_client_module() Widgets = _WidgetsEnumType("Widgets", {name: name for name in _plugin_widgets} | _Widgets) diff --git a/bec_widgets/cli/rpc/rpc_widget_handler.py b/bec_widgets/cli/rpc/rpc_widget_handler.py index 669db3fd..b261d1b1 100644 --- a/bec_widgets/cli/rpc/rpc_widget_handler.py +++ b/bec_widgets/cli/rpc/rpc_widget_handler.py @@ -31,10 +31,9 @@ class RPCWidgetHandler: Returns: None """ - clss = get_custom_classes("bec_widgets") - self._widget_classes = get_all_plugin_widgets() | { - cls.__name__: cls for cls in clss.widgets if cls.__name__ not in IGNORE_WIDGETS - } + self._widget_classes = ( + get_custom_classes("bec_widgets") + get_all_plugin_widgets() + ).as_dict(IGNORE_WIDGETS) def create_widget(self, widget_type, **kwargs) -> BECWidget: """ diff --git a/bec_widgets/utils/bec_plugin_helper.py b/bec_widgets/utils/bec_plugin_helper.py index d9641af8..8f5ca0c8 100644 --- a/bec_widgets/utils/bec_plugin_helper.py +++ b/bec_widgets/utils/bec_plugin_helper.py @@ -3,12 +3,17 @@ from __future__ import annotations import importlib.metadata import inspect import pkgutil +import traceback from importlib import util as importlib_util from importlib.machinery import FileFinder, ModuleSpec, SourceFileLoader from types import ModuleType from typing import Generator -from bec_widgets.utils.bec_widget import BECWidget +from bec_lib.logger import bec_logger + +from bec_widgets.utils.plugin_utils import BECClassContainer, BECClassInfo + +logger = bec_logger.logger def _submodule_specs(module: ModuleType) -> tuple[ModuleSpec | None, ...]: @@ -30,7 +35,12 @@ def _loaded_submodules_from_specs( assert isinstance( submodule.__loader__, SourceFileLoader ), "Module found from FileFinder should have SourceFileLoader!" - submodule.__loader__.exec_module(submodule) + try: + submodule.__loader__.exec_module(submodule) + except Exception as e: + logger.error( + f"Error loading plugin {submodule}: \n{''.join(traceback.format_exception(e))}" + ) yield submodule @@ -41,27 +51,29 @@ def _submodule_by_name(module: ModuleType, name: str): return None -def _get_widgets_from_module(module: ModuleType) -> dict[str, "type[BECWidget]"]: - """Find any BECWidget subclasses in the given module and return them with their names.""" +def _get_widgets_from_module(module: ModuleType) -> BECClassContainer: + """Find any BECWidget subclasses in the given module and return them with their info.""" from bec_widgets.utils.bec_widget import BECWidget # avoid circular import - return dict( - inspect.getmembers( - module, - predicate=lambda item: inspect.isclass(item) - and issubclass(item, BECWidget) - and item is not BECWidget, - ) + classes = inspect.getmembers( + module, + predicate=lambda item: inspect.isclass(item) + and issubclass(item, BECWidget) + and item is not BECWidget, + ) + return BECClassContainer( + BECClassInfo(name=k, module=module.__name__, file=module.__loader__.get_filename(), obj=v) + for k, v in classes ) -def _all_widgets_from_all_submods(module): +def _all_widgets_from_all_submods(module) -> BECClassContainer: """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 for submod in _loaded_submodules_from_specs(_submodule_specs(module)): - widgets.update(_all_widgets_from_all_submods(submod)) + widgets += _all_widgets_from_all_submods(submod) return widgets @@ -75,15 +87,16 @@ def get_plugin_client_module() -> ModuleType | None: return _submodule_by_name(plugin, "client") if (plugin := user_widget_plugin()) else None -def get_all_plugin_widgets() -> dict[str, "type[BECWidget]"]: +def get_all_plugin_widgets() -> BECClassContainer: """If there is a plugin repository installed, load all widgets from it.""" if plugin := user_widget_plugin(): return _all_widgets_from_all_submods(plugin) else: - return {} + return BECClassContainer() if __name__ == "__main__": # pragma: no cover - # print(get_all_plugin_widgets()) + client = get_plugin_client_module() + print(get_all_plugin_widgets()) ... diff --git a/bec_widgets/utils/plugin_utils.py b/bec_widgets/utils/plugin_utils.py index 7b6d4f6a..434abf59 100644 --- a/bec_widgets/utils/plugin_utils.py +++ b/bec_widgets/utils/plugin_utils.py @@ -4,7 +4,7 @@ import importlib import inspect import os from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Iterable from bec_lib.plugin_helper import _get_available_plugins from qtpy.QtWidgets import QGraphicsWidget, QWidget @@ -90,15 +90,15 @@ class BECClassInfo: name: str module: str file: str - obj: type + obj: type[BECWidget] is_connector: bool = False is_widget: bool = False is_plugin: bool = False class BECClassContainer: - def __init__(self): - self._collection: list[BECClassInfo] = [] + def __init__(self, initial: Iterable[BECClassInfo] = []): + self._collection: list[BECClassInfo] = list(initial) def __repr__(self): return str(list(cl.name for cl in self.collection)) @@ -106,6 +106,16 @@ class BECClassContainer: def __iter__(self): return self._collection.__iter__() + def __add__(self, other: BECClassContainer): + return BECClassContainer((*self, *(c for c in other if c.name not in self.names))) + + def as_dict(self, ignores: list[str] = []) -> dict[str, type[BECWidget]]: + """get a dict of {name: Type} for all the entries in the collection. + + Args: + ignores(list[str]): a list of class names to exclude from the dictionary.""" + return {c.name: c.obj for c in self if c.name not in ignores} + def add_class(self, class_info: BECClassInfo): """ Add a class to the collection. @@ -115,53 +125,44 @@ class BECClassContainer: """ self.collection.append(class_info) + @property + def names(self): + """Return a list of class names""" + return [c.name for c in self] + @property def collection(self): - """ - Get the collection of classes. - """ + """Get the collection of classes.""" return self._collection @property def connector_classes(self): - """ - Get all connector classes. - """ + """Get all connector classes.""" return [info.obj for info in self.collection if info.is_connector] @property def top_level_classes(self): - """ - Get all top-level classes. - """ + """Get all top-level classes.""" return [info.obj for info in self.collection if info.is_plugin] @property def plugins(self): - """ - Get all plugins. These are all classes that are on the top level and are widgets. - """ + """Get all plugins. These are all classes that are on the top level and are widgets.""" return [info.obj for info in self.collection if info.is_widget and info.is_plugin] @property def widgets(self): - """ - Get all widgets. These are all classes inheriting from BECWidget. - """ + """Get all widgets. These are all classes inheriting from BECWidget.""" return [info.obj for info in self.collection if info.is_widget] @property def rpc_top_level_classes(self): - """ - Get all top-level classes that are RPC-enabled. These are all classes that users can choose from. - """ + """Get all top-level classes that are RPC-enabled. These are all classes that users can choose from.""" return [info.obj for info in self.collection if info.is_plugin and info.is_connector] @property def classes(self): - """ - Get all classes. - """ + """Get all classes.""" return [info.obj for info in self.collection] @@ -197,7 +198,7 @@ def get_custom_classes(repo_name: str) -> BECClassContainer: if not hasattr(obj, "__module__") or obj.__module__ != module.__name__: continue if isinstance(obj, type): - class_info = BECClassInfo(name=name, module=module_name, file=path, obj=obj) + class_info = BECClassInfo(name=name, module=module.__name__, file=path, obj=obj) if issubclass(obj, BECConnector): class_info.is_connector = True if issubclass(obj, BECWidget): diff --git a/bec_widgets/utils/ui_loader.py b/bec_widgets/utils/ui_loader.py index 76251ce5..594bd76f 100644 --- a/bec_widgets/utils/ui_loader.py +++ b/bec_widgets/utils/ui_loader.py @@ -31,12 +31,9 @@ class UILoader: def __init__(self, parent=None): self.parent = parent - widgets = get_custom_classes("bec_widgets").classes - - self.custom_widgets = {widget.__name__: widget for widget in widgets} - - plugin_widgets = get_all_plugin_widgets() - self.custom_widgets.update(plugin_widgets) + self.custom_widgets = ( + get_custom_classes("bec_widgets") + get_all_plugin_widgets() + ).as_dict() if PYSIDE6: self.loader = self.load_ui_pyside6 diff --git a/tests/unit_tests/test_bec_plugin_helper.py b/tests/unit_tests/test_bec_plugin_helper.py index 94e8ff35..522814a8 100644 --- a/tests/unit_tests/test_bec_plugin_helper.py +++ b/tests/unit_tests/test_bec_plugin_helper.py @@ -2,7 +2,9 @@ from importlib.machinery import FileFinder, SourceFileLoader from types import ModuleType from unittest import mock -from bec_widgets.utils.bec_plugin_helper import BECWidget, _all_widgets_from_all_submods +from bec_widgets.utils.bec_plugin_helper import _all_widgets_from_all_submods +from bec_widgets.utils.bec_widget import BECWidget +from bec_widgets.utils.plugin_utils import BECClassContainer, BECClassInfo def test_all_widgets_from_module_no_submodules(): @@ -39,10 +41,17 @@ def test_all_widgets_from_module_with_submodules(): mock.patch("importlib.util.module_from_spec", return_value=submodule), mock.patch( "bec_widgets.utils.bec_plugin_helper._get_widgets_from_module", - side_effect=[{"TestWidget": BECWidget}, {"SubWidget": BECWidget}], + side_effect=[ + BECClassContainer( + [BECClassInfo(name="TestWidget", module="", obj=BECWidget, file="")] + ), + BECClassContainer( + [BECClassInfo(name="SubWidget", module="", obj=BECWidget, file="")] + ), + ], ), ): - widgets = _all_widgets_from_all_submods(module) + widgets = _all_widgets_from_all_submods(module).as_dict() assert widgets == {"TestWidget": BECWidget, "SubWidget": BECWidget} @@ -54,8 +63,9 @@ def test_all_widgets_from_module_no_widgets(): module = mock.MagicMock() with mock.patch( - "bec_widgets.utils.bec_plugin_helper._get_widgets_from_module", return_value={} + "bec_widgets.utils.bec_plugin_helper._get_widgets_from_module", + return_value=BECClassContainer([]), ): - widgets = _all_widgets_from_all_submods(module) + widgets = _all_widgets_from_all_submods(module).as_dict() assert widgets == {} diff --git a/tests/unit_tests/test_client_plugin_widgets.py b/tests/unit_tests/test_client_plugin_widgets.py index 4bafb15a..b43c2960 100644 --- a/tests/unit_tests/test_client_plugin_widgets.py +++ b/tests/unit_tests/test_client_plugin_widgets.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, call, patch from bec_widgets.cli import client from bec_widgets.cli.rpc.rpc_base import RPCBase +from bec_widgets.utils.plugin_utils import BECClassContainer, BECClassInfo class _TestGlobalPlugin(RPCBase): ... @@ -47,7 +48,9 @@ mock_client_module_duplicate.DeviceComboBox = _TestDuplicatePlugin ) @patch( "bec_widgets.utils.bec_plugin_helper.get_all_plugin_widgets", - return_value={"DeviceComboBox": _TestDuplicatePlugin}, + return_value=BECClassContainer( + [BECClassInfo(name="DeviceComboBox", obj=_TestDuplicatePlugin, module="", file="")] + ), ) def test_duplicate_plugins_not_allowed(_, bec_logger: MagicMock): reload(client) diff --git a/tests/unit_tests/test_generate_cli_client.py b/tests/unit_tests/test_generate_cli_client.py index 2ecf3915..09e0fe72 100644 --- a/tests/unit_tests/test_generate_cli_client.py +++ b/tests/unit_tests/test_generate_cli_client.py @@ -99,7 +99,7 @@ def test_client_generator_with_black_formatting(): try: - _plugin_widgets = get_all_plugin_widgets() + _plugin_widgets = get_all_plugin_widgets().as_dict() plugin_client = get_plugin_client_module() Widgets = _WidgetsEnumType("Widgets", {name: name for name in _plugin_widgets} | _Widgets) diff --git a/tests/unit_tests/test_rpc_widget_handler.py b/tests/unit_tests/test_rpc_widget_handler.py index cbdea1a8..1f2fc768 100644 --- a/tests/unit_tests/test_rpc_widget_handler.py +++ b/tests/unit_tests/test_rpc_widget_handler.py @@ -7,6 +7,7 @@ from bec_widgets.cli import client from bec_widgets.cli.rpc.rpc_base import RPCBase from bec_widgets.cli.rpc.rpc_widget_handler import RPCWidgetHandler from bec_widgets.utils.bec_widget import BECWidget +from bec_widgets.utils.plugin_utils import BECClassContainer, BECClassInfo from bec_widgets.widgets.containers.dock.dock import BECDock @@ -21,7 +22,12 @@ class _TestPluginWidget(BECWidget): ... @patch( "bec_widgets.cli.rpc.rpc_widget_handler.get_all_plugin_widgets", - return_value={"DeviceComboBox": _TestPluginWidget, "NewPluginWidget": _TestPluginWidget}, + return_value=BECClassContainer( + [ + BECClassInfo(name="DeviceComboBox", obj=_TestPluginWidget, module="", file=""), + BECClassInfo(name="NewPluginWidget", obj=_TestPluginWidget, module="", file=""), + ] + ), ) def test_duplicate_plugins_not_allowed(_): handler = RPCWidgetHandler() diff --git a/tests/unit_tests/test_scan_control.py b/tests/unit_tests/test_scan_control.py index d47f661b..9696809a 100644 --- a/tests/unit_tests/test_scan_control.py +++ b/tests/unit_tests/test_scan_control.py @@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch import pytest from bec_lib.endpoints import MessageEndpoints from bec_lib.messages import AvailableResourceMessage, ScanQueueHistoryMessage, ScanQueueMessage -from qtpy.QtCore import QModelIndex, QPoint, Qt +from qtpy.QtCore import QModelIndex, Qt from bec_widgets.utils.forms_from_types.items import StrMetadataField from bec_widgets.utils.widget_io import WidgetIO