From 152aadfffd79ecb28f40149f532bbd97b60f1edf Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Wed, 15 Apr 2026 10:48:17 +0200 Subject: [PATCH] perf(rpc): build widget registry from lazy references --- .../device_manager_display_widget.py | 34 ++- .../device_manager_view.py | 9 +- .../device_manager_widget.py | 9 +- bec_widgets/cli/rpc/rpc_widget_handler.py | 42 +++- bec_widgets/utils/bec_plugin_helper.py | 62 ++++- bec_widgets/utils/plugin_utils.py | 237 ++++++++++++++---- 6 files changed, 315 insertions(+), 78 deletions(-) diff --git a/bec_widgets/applications/views/device_manager_view/device_manager_display_widget.py b/bec_widgets/applications/views/device_manager_view/device_manager_display_widget.py index 248f50db..d6915d2e 100644 --- a/bec_widgets/applications/views/device_manager_view/device_manager_display_widget.py +++ b/bec_widgets/applications/views/device_manager_view/device_manager_display_widget.py @@ -13,6 +13,7 @@ from bec_lib.file_utils import DeviceConfigWriter from bec_lib.logger import bec_logger from bec_lib.messages import ConfigAction, ScanStatusMessage from bec_lib.plugin_helper import plugin_package_name, plugin_repo_path +from bec_lib.utils.import_utils import lazy_import_from from bec_qthemes import apply_theme, material_icon from qtpy.QtCore import QMetaObject, Qt, QThreadPool, Signal from qtpy.QtGui import QColor @@ -26,15 +27,6 @@ from qtpy.QtWidgets import ( QWidget, ) -from bec_widgets.applications.views.device_manager_view.device_manager_dialogs.config_choice_dialog import ( - ConfigChoiceDialog, -) -from bec_widgets.applications.views.device_manager_view.device_manager_dialogs.device_form_dialog import ( - DeviceFormDialog, -) -from bec_widgets.applications.views.device_manager_view.device_manager_dialogs.upload_redis_dialog import ( - UploadRedisDialog, -) from bec_widgets.utils.colors import get_accent_colors from bec_widgets.utils.error_popups import SafeSlot from bec_widgets.utils.toolbars.actions import MaterialIconAction @@ -47,9 +39,6 @@ from bec_widgets.widgets.control.device_manager.components.device_table.device_t ) from bec_widgets.widgets.control.device_manager.components.dm_config_view import DMConfigView from bec_widgets.widgets.control.device_manager.components.dm_docstring_view import DocstringView -from bec_widgets.widgets.control.device_manager.components.ophyd_validation.ophyd_validation import ( - OphydValidation, -) from bec_widgets.widgets.control.device_manager.components.ophyd_validation.ophyd_validation_utils import ( ConfigStatus, ConnectionStatus, @@ -65,8 +54,29 @@ from bec_widgets.widgets.utility.spinner.spinner import SpinnerWidget if TYPE_CHECKING: # pragma: no cover from bec_lib.client import BECClient + from bec_widgets.applications.views.device_manager_view.device_manager_dialogs.upload_redis_dialog import ( + UploadRedisDialog, + ) + logger = bec_logger.logger +ConfigChoiceDialog = lazy_import_from( + "bec_widgets.applications.views.device_manager_view.device_manager_dialogs.config_choice_dialog", + ("ConfigChoiceDialog",), +) +DeviceFormDialog = lazy_import_from( + "bec_widgets.applications.views.device_manager_view.device_manager_dialogs.device_form_dialog", + ("DeviceFormDialog",), +) +UploadRedisDialog = lazy_import_from( + "bec_widgets.applications.views.device_manager_view.device_manager_dialogs.upload_redis_dialog", + ("UploadRedisDialog",), +) +OphydValidation = lazy_import_from( + "bec_widgets.widgets.control.device_manager.components.ophyd_validation.ophyd_validation", + ("OphydValidation",), +) + _yes_no_question = partial( QMessageBox.question, buttons=QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, diff --git a/bec_widgets/applications/views/device_manager_view/device_manager_view.py b/bec_widgets/applications/views/device_manager_view/device_manager_view.py index 79770847..36491fca 100644 --- a/bec_widgets/applications/views/device_manager_view/device_manager_view.py +++ b/bec_widgets/applications/views/device_manager_view/device_manager_view.py @@ -1,14 +1,17 @@ """Module for Device Manager View.""" +from bec_lib.utils.import_utils import lazy_import_from from qtpy.QtCore import QRect from qtpy.QtWidgets import QWidget -from bec_widgets.applications.views.device_manager_view.device_manager_widget import ( - DeviceManagerWidget, -) from bec_widgets.applications.views.view import ViewBase, ViewTourSteps from bec_widgets.utils.error_popups import SafeSlot +DeviceManagerWidget = lazy_import_from( + "bec_widgets.applications.views.device_manager_view.device_manager_widget", + ("DeviceManagerWidget",), +) + class DeviceManagerView(ViewBase): """ diff --git a/bec_widgets/applications/views/device_manager_view/device_manager_widget.py b/bec_widgets/applications/views/device_manager_view/device_manager_widget.py index d6201500..0658abcf 100644 --- a/bec_widgets/applications/views/device_manager_view/device_manager_widget.py +++ b/bec_widgets/applications/views/device_manager_view/device_manager_widget.py @@ -6,15 +6,18 @@ import os from bec_lib.bec_yaml_loader import yaml_load from bec_lib.logger import bec_logger +from bec_lib.utils.import_utils import lazy_import_from from bec_qthemes import material_icon from qtpy import QtCore, QtWidgets -from bec_widgets.applications.views.device_manager_view.device_manager_display_widget import ( - DeviceManagerDisplayWidget, -) from bec_widgets.utils.bec_widget import BECWidget from bec_widgets.utils.error_popups import SafeSlot +DeviceManagerDisplayWidget = lazy_import_from( + "bec_widgets.applications.views.device_manager_view.device_manager_display_widget", + ("DeviceManagerDisplayWidget",), +) + logger = bec_logger.logger diff --git a/bec_widgets/cli/rpc/rpc_widget_handler.py b/bec_widgets/cli/rpc/rpc_widget_handler.py index 83d9a04d..625fc636 100644 --- a/bec_widgets/cli/rpc/rpc_widget_handler.py +++ b/bec_widgets/cli/rpc/rpc_widget_handler.py @@ -1,9 +1,19 @@ from __future__ import annotations -from bec_widgets.cli.client_utils import IGNORE_WIDGETS -from bec_widgets.utils.bec_plugin_helper import get_all_plugin_widgets -from bec_widgets.utils.bec_widget import BECWidget -from bec_widgets.utils.plugin_utils import get_custom_classes +from typing import TYPE_CHECKING + +from bec_lib.utils.import_utils import lazy_import_from + +from bec_widgets.utils.bec_plugin_helper import get_all_plugin_widget_references +from bec_widgets.utils.plugin_utils import get_custom_class_references + +try: + from bec_widgets.cli.constants import IGNORE_WIDGETS +except ModuleNotFoundError: # pragma: no cover + IGNORE_WIDGETS = ["LaunchWindow"] + +if TYPE_CHECKING: # pragma: no cover + from bec_widgets.utils.bec_widget import BECWidget class RPCWidgetHandler: @@ -13,7 +23,7 @@ class RPCWidgetHandler: self._widget_classes = None @property - def widget_classes(self) -> dict[str, type[BECWidget]]: + def widget_classes(self) -> dict[str, type["BECWidget"]]: """ Get the available widget classes. @@ -31,12 +41,24 @@ class RPCWidgetHandler: Returns: None """ - self._widget_classes = ( - get_custom_classes("bec_widgets", packages=("widgets", "applications")) - + get_all_plugin_widgets() - ).as_dict(IGNORE_WIDGETS) + ignored = set(IGNORE_WIDGETS) + widget_classes = { + reference.name: lazy_import_from(reference.module, (reference.name,)) + for reference in get_all_plugin_widget_references(use_cache=False) + if reference.name not in ignored + } + widget_classes.update( + { + reference.name: lazy_import_from(reference.module, (reference.name,)) + for reference in get_custom_class_references( + "bec_widgets", packages=("widgets", "applications"), use_cache=False + ) + if reference.name not in ignored + } + ) + self._widget_classes = widget_classes - def create_widget(self, widget_type, **kwargs) -> BECWidget: + def create_widget(self, widget_type, **kwargs) -> "BECWidget": """ Create a widget from an RPC message. diff --git a/bec_widgets/utils/bec_plugin_helper.py b/bec_widgets/utils/bec_plugin_helper.py index 242adbdc..c34a04a0 100644 --- a/bec_widgets/utils/bec_plugin_helper.py +++ b/bec_widgets/utils/bec_plugin_helper.py @@ -1,7 +1,9 @@ from __future__ import annotations +import ast import importlib.metadata import inspect +import logging import pkgutil import traceback from importlib import util as importlib_util @@ -9,11 +11,65 @@ from importlib.machinery import FileFinder, ModuleSpec, SourceFileLoader from types import ModuleType from typing import Generator -from bec_lib.logger import bec_logger +from bec_widgets.utils.plugin_utils import ( + BECClassContainer, + BECClassInfo, + BECClassReference, + _ast_node_name, + _class_has_rpc_markers, + _discover_class_references_from_roots, + _find_package_roots, +) -from bec_widgets.utils.plugin_utils import BECClassContainer, BECClassInfo +logger = logging.getLogger(__name__) -logger = bec_logger.logger + +def _plugin_class_is_candidate(node: ast.ClassDef) -> bool: + base_names = {_ast_node_name(base) for base in node.bases} + return bool({"BECWidget", "BECConnector"} & base_names) or _class_has_rpc_markers(node) + + +_PLUGIN_WIDGET_REFERENCE_CACHE: dict[tuple[tuple[str, str], ...], tuple[BECClassReference, ...]] = ( + {} +) + + +def _plugin_entry_point_snapshot() -> tuple[tuple[str, str], ...]: + return tuple( + sorted( + (entry_point.name, entry_point.module) + for entry_point in importlib.metadata.entry_points(group="bec.widgets.user_widgets") # type: ignore + ) + ) + + +def _build_plugin_widget_references() -> tuple[BECClassReference, ...]: + references: list[BECClassReference] = [] + seen_names: set[str] = set() + for entry_point in importlib.metadata.entry_points(group="bec.widgets.user_widgets"): # type: ignore + try: + package_roots = _find_package_roots(entry_point.module) + except ModuleNotFoundError: + continue + for reference in _discover_class_references_from_roots( + entry_point.module, + package_roots, + file_name_filter=lambda file_name: file_name.endswith(".py") + and not file_name.startswith("__"), + candidate_filter=_plugin_class_is_candidate, + ): + if reference.name in seen_names: + continue + references.append(reference) + seen_names.add(reference.name) + return tuple(references) + + +def get_all_plugin_widget_references(*, use_cache: bool = True) -> list[BECClassReference]: + snapshot = _plugin_entry_point_snapshot() + if not use_cache or snapshot not in _PLUGIN_WIDGET_REFERENCE_CACHE: + _PLUGIN_WIDGET_REFERENCE_CACHE[snapshot] = _build_plugin_widget_references() + return list(_PLUGIN_WIDGET_REFERENCE_CACHE[snapshot]) def _submodule_specs(module: ModuleType) -> tuple[ModuleSpec | None, ...]: diff --git a/bec_widgets/utils/plugin_utils.py b/bec_widgets/utils/plugin_utils.py index c9c367ae..bf672e6c 100644 --- a/bec_widgets/utils/plugin_utils.py +++ b/bec_widgets/utils/plugin_utils.py @@ -1,22 +1,23 @@ from __future__ import annotations +import ast import importlib import inspect import os from dataclasses import dataclass -from typing import TYPE_CHECKING, Iterable - -from bec_lib.plugin_helper import _get_available_plugins -from qtpy.QtWidgets import QWidget - -from bec_widgets.utils.bec_connector import BECConnector -from bec_widgets.utils.bec_widget import BECWidget +from functools import lru_cache +from importlib import util as importlib_util +from typing import TYPE_CHECKING, Callable, Iterable if TYPE_CHECKING: # pragma: no cover + from bec_widgets.utils.bec_connector import BECConnector + from bec_widgets.utils.bec_widget import BECWidget from bec_widgets.widgets.containers.auto_update.auto_updates import AutoUpdates +_DISCOVERY_BASE_NAMES = frozenset({"BECConnector", "BECWidget", "ViewBase"}) -def get_plugin_widgets() -> dict[str, BECConnector]: + +def get_plugin_widgets() -> dict[str, "BECConnector"]: """ Get all available widgets from the plugin directory. Widgets are classes that inherit from BECConnector. The plugins are provided through python plugins and specified in the respective pyproject.toml file using @@ -35,9 +36,10 @@ def get_plugin_widgets() -> dict[str, BECConnector]: Returns: dict[str, BECConnector]: A dictionary of widget names and their respective classes. """ + from bec_lib.plugin_helper import _get_available_plugins + modules = _get_available_plugins("bec.widgets.user_widgets") loaded_plugins = {} - print(modules) for module in modules: mods = inspect.getmembers(module, predicate=_filter_plugins) for name, mod_cls in mods: @@ -48,6 +50,8 @@ def get_plugin_widgets() -> dict[str, BECConnector]: def _filter_plugins(obj): + from bec_widgets.utils.bec_connector import BECConnector + return inspect.isclass(obj) and issubclass(obj, BECConnector) @@ -66,6 +70,8 @@ def get_plugin_auto_updates() -> dict[str, type[AutoUpdates]]: Returns: dict[str, AutoUpdates]: A dictionary of widget names and their respective classes. """ + from bec_lib.plugin_helper import _get_available_plugins + modules = _get_available_plugins("bec.widgets.auto_updates") loaded_plugins = {} for module in modules: @@ -90,14 +96,20 @@ class BECClassInfo: name: str module: str file: str - obj: type[BECWidget] + obj: type["BECWidget"] is_connector: bool = False is_widget: bool = False is_plugin: bool = False +@dataclass(frozen=True) +class BECClassReference: + name: str + module: str + + class BECClassContainer: - def __init__(self, initial: Iterable[BECClassInfo] = []): + def __init__(self, initial: Iterable[BECClassInfo] = ()): self._collection: list[BECClassInfo] = list(initial) def __repr__(self): @@ -109,12 +121,13 @@ class BECClassContainer: 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]]: + def as_dict(self, ignores: list[str] | None = None) -> 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} + ignore_set = set(ignores or ()) + return {c.name: c.obj for c in self if c.name not in ignore_set} def add_class(self, class_info: BECClassInfo): """ @@ -166,48 +179,178 @@ class BECClassContainer: return [info.obj for info in self.collection] +def _ast_node_name(node: ast.expr) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return None + + +def _class_has_rpc_markers(node: ast.ClassDef) -> bool: + for stmt in node.body: + if isinstance(stmt, ast.Assign): + target_names = {target.id for target in stmt.targets if isinstance(target, ast.Name)} + if ( + "PLUGIN" in target_names + and isinstance(stmt.value, ast.Constant) + and stmt.value.value + ): + return True + if {"RPC_CONTENT_CLASS", "USER_ACCESS"} & target_names: + return True + if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): + if ( + stmt.target.id == "PLUGIN" + and isinstance(stmt.value, ast.Constant) + and stmt.value.value + ): + return True + if stmt.target.id in {"RPC_CONTENT_CLASS", "USER_ACCESS"}: + return True + return False + + +def _class_is_candidate(node: ast.ClassDef) -> bool: + base_names = {_ast_node_name(base) for base in node.bases} + return bool(_DISCOVERY_BASE_NAMES & base_names) or _class_has_rpc_markers(node) + + +def _candidate_top_level_class_names(path: str) -> list[str]: + with open(path, encoding="utf-8") as file_handle: + module = ast.parse(file_handle.read(), filename=path) + return [ + node.name + for node in module.body + if isinstance(node, ast.ClassDef) and _class_is_candidate(node) + ] + + +@lru_cache(maxsize=64) +def _find_package_roots(module_name: str) -> tuple[str, ...]: + spec = importlib_util.find_spec(module_name) + if spec is None: + raise ModuleNotFoundError(module_name) + + package_roots = tuple(spec.submodule_search_locations or ()) + if package_roots: + return package_roots + if spec.origin: + return (os.path.dirname(spec.origin),) + raise ModuleNotFoundError(module_name) + + +def _discover_class_references_from_roots( + module_prefix: str, + package_roots: Iterable[str], + *, + file_name_filter: Callable[[str], bool], + candidate_filter: Callable[[ast.ClassDef], bool], +) -> tuple[BECClassReference, ...]: + references: list[BECClassReference] = [] + seen_names: set[str] = set() + + for package_root in package_roots: + for root, _, files in sorted(os.walk(package_root)): + for file_name in sorted(files): + if not file_name_filter(file_name): + continue + path = os.path.join(root, file_name) + with open(path, encoding="utf-8") as file_handle: + module = ast.parse(file_handle.read(), filename=path) + rel_path = os.path.relpath(path, package_root).removesuffix(".py") + module_name = ".".join([module_prefix, *rel_path.split(os.sep)]) + for node in module.body: + if not isinstance(node, ast.ClassDef) or not candidate_filter(node): + continue + if node.name in seen_names: + continue + references.append(BECClassReference(name=node.name, module=module_name)) + seen_names.add(node.name) + + return tuple(references) + + +def _iter_candidate_modules(repo_name: str, package: str) -> Iterable[tuple[str, str, list[str]]]: + try: + package_roots = _find_package_roots(f"{repo_name}.{package}") + except ModuleNotFoundError: + return () + + modules: list[tuple[str, str, list[str]]] = [] + for directory in package_roots: + for root, _, files in sorted(os.walk(directory)): + for file_name in sorted(files): + if ( + not file_name.endswith(".py") + or file_name.startswith("__") + or file_name.startswith("register_") + or file_name.endswith("_plugin.py") + ): + continue + path = os.path.join(root, file_name) + rel_dir = os.path.dirname(os.path.relpath(path, directory)) + module_name = ( + file_name.removesuffix(".py") + if rel_dir in ("", ".") + else ".".join(rel_dir.split(os.sep) + [file_name.removesuffix(".py")]) + ) + class_names = _candidate_top_level_class_names(path) + if class_names: + modules.append((f"{repo_name}.{package}.{module_name}", path, class_names)) + return tuple(modules) + + def _collect_classes_from_package(repo_name: str, package: str) -> BECClassContainer: """Collect classes from a package subtree (for example ``widgets`` or ``applications``).""" collection = BECClassContainer() - try: - anchor_module = importlib.import_module(f"{repo_name}.{package}") - except ModuleNotFoundError as exc: - # Some plugin repositories expose only one subtree. Skip gracefully if it does not exist. - if exc.name == f"{repo_name}.{package}": - return collection - raise + for module_name, path, _ in _iter_candidate_modules(repo_name, package): + from qtpy.QtWidgets import QWidget - directory = os.path.dirname(anchor_module.__file__) - for root, _, files in sorted(os.walk(directory)): - for file in files: - if not file.endswith(".py") or file.startswith("__"): + from bec_widgets.utils.bec_connector import BECConnector + from bec_widgets.utils.bec_widget import BECWidget + + module = importlib.import_module(module_name) + for name, obj in inspect.getmembers(module, inspect.isclass): + if obj.__module__ != module.__name__: continue - - path = os.path.join(root, file) - rel_dir = os.path.dirname(os.path.relpath(path, directory)) - if rel_dir in ("", "."): - module_name = file.split(".")[0] - else: - module_name = ".".join(rel_dir.split(os.sep) + [file.split(".")[0]]) - - module = importlib.import_module(f"{repo_name}.{package}.{module_name}") - - for name in dir(module): - obj = getattr(module, name) - 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) - if issubclass(obj, BECConnector): - class_info.is_connector = True - if issubclass(obj, QWidget) or issubclass(obj, BECWidget): - class_info.is_widget = True - if hasattr(obj, "PLUGIN") and obj.PLUGIN: - class_info.is_plugin = True - collection.add_class(class_info) + class_info = BECClassInfo(name=name, module=module.__name__, file=path, obj=obj) + if issubclass(obj, BECConnector): + class_info.is_connector = True + if issubclass(obj, QWidget) or issubclass(obj, BECWidget): + class_info.is_widget = True + if hasattr(obj, "PLUGIN") and obj.PLUGIN: + class_info.is_plugin = True + collection.add_class(class_info) return collection +@lru_cache(maxsize=32) +def _cached_custom_class_references( + repo_name: str, packages: tuple[str, ...] +) -> tuple[BECClassReference, ...]: + references: list[BECClassReference] = [] + seen_names: set[str] = set() + for package in packages: + for module_name, _, class_names in _iter_candidate_modules(repo_name, package): + for class_name in class_names: + if class_name in seen_names: + continue + references.append(BECClassReference(name=class_name, module=module_name)) + seen_names.add(class_name) + return tuple(references) + + +def get_custom_class_references( + repo_name: str, packages: tuple[str, ...] | None = None, *, use_cache: bool = True +) -> list[BECClassReference]: + selected_packages = packages or ("widgets",) + if use_cache: + return list(_cached_custom_class_references(repo_name, tuple(selected_packages))) + _cached_custom_class_references.cache_clear() + return list(_cached_custom_class_references(repo_name, tuple(selected_packages))) + + def get_custom_classes( repo_name: str, packages: tuple[str, ...] | None = None ) -> BECClassContainer: