perf(rpc): build widget registry from lazy references

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,

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):
     """

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
 
 

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.
 

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, ...]:

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: