feat: add loader/helper for widget plugins

bec_widgets/cli/generate_cli.py

@@ -34,13 +34,25 @@ else:
 
 
 class ClientGenerator:
-    def __init__(self):
-        self.header = """# This file was automatically generated by generate_cli.py\n
+    def __init__(self, base=False):
+        self._base = base
+        base_imports = (
+            """import enum
+import inspect
+from typing import Literal, Optional
+"""
+            if self._base
+            else "\n"
+        )
+        self.header = f"""# This file was automatically generated by generate_cli.py\n
 from __future__ import annotations
-import enum
-from typing import Literal, Optional, overload
+{base_imports}
+from bec_lib.logger import bec_logger
 
 from bec_widgets.cli.rpc.rpc_base import RPCBase, rpc_call
+{"from bec_widgets.utils.bec_plugin_helper import get_all_plugin_widgets, get_plugin_client_module" if self._base else ""}
+
+logger = bec_logger.logger
 
 # pylint: skip-file"""
 
@@ -67,6 +79,7 @@ from bec_widgets.cli.rpc.rpc_base import RPCBase, rpc_call
 
         self.write_client_enum(rpc_top_level_classes)
         for cls in connector_classes:
+            logger.debug(f"generating RPC client class for {cls.__name__}")
             self.content += "\n\n"
             self.generate_content_for_class(cls)
 
@@ -74,10 +87,14 @@ from bec_widgets.cli.rpc.rpc_base import RPCBase, rpc_call
         """
         Write the client enum to the content.
         """
-        self.content += """
+        if self._base:
+            self.content += """
 class _WidgetsEnumType(str, enum.Enum):
     \"\"\" Enum for the available widgets, to be generated programatically  \"\"\"
     ...
+    """
+
+        self.content += """
 
 _Widgets = {
     """
@@ -85,8 +102,33 @@ _Widgets = {
             self.content += f'"{cls.__name__}": "{cls.__name__}",\n    '
 
         self.content += """}
-Widgets = _WidgetsEnumType("Widgets", _Widgets)
-        """
+"""
+        if self._base:
+            self.content += """
+
+_plugin_widgets = get_all_plugin_widgets()
+plugin_client = get_plugin_client_module()
+Widgets = _WidgetsEnumType("Widgets", {name: name for name in _plugin_widgets} | _Widgets)
+
+if (_overlap := _Widgets.keys() & _plugin_widgets.keys()) != set():
+    for _widget in _overlap:
+        logger.warning(f"Detected duplicate widget {_widget} in plugin repo file: {inspect.getfile(_plugin_widgets[_widget])} !")
+for plugin_name, plugin_class in inspect.getmembers(plugin_client, inspect.isclass):
+    if issubclass(plugin_class, RPCBase) and plugin_class is not RPCBase:
+        if plugin_name in globals():
+            conflicting_file = (
+                inspect.getfile(_plugin_widgets[plugin_name])
+                if plugin_name in _plugin_widgets
+                else f"{plugin_client}"
+            )
+            logger.warning(
+                f"Plugin widget {plugin_name} from {conflicting_file} conflicts with a built-in class!"
+            )
+            continue
+        if plugin_name not in _overlap:
+            globals()[plugin_name] = plugin_class
+
+"""
 
     def generate_content_for_class(self, cls):
         """
@@ -199,38 +241,59 @@ def main():
 
     parser = argparse.ArgumentParser(description="Auto-generate the client for RPC widgets")
     parser.add_argument(
-        "--module-name",
+        "--target",
         action="store",
         type=str,
-        default="bec_widgets",
-        help="Which module to generate plugin files for (default: bec_widgets, example: my_plugin_repo.bec_widgets)",
+        help="Which package to generate plugin files for. Should be installed in the local environment (example: my_plugin_repo)",
     )
 
     args = parser.parse_args()
+    if args.target is None:
+        logger.error(
+            "You must provide a target - for safety, the default of running this on bec_widgets core has been removed. To generate the client for bec_widgets, run `bw-generate-cli --target bec_widgets`"
+        )
+        return
 
     logger.info(f"BEC Widget code generation tool started with args: {args}")
 
+    client_subdir = "cli" if args.target == "bec_widgets" else "widgets"
+    module_name = "bec_widgets" if args.target == "bec_widgets" else f"{args.target}.bec_widgets"
+
     try:
-        module = importlib.import_module(args.module_name)
+        module = importlib.import_module(module_name)
         assert module.__file__ is not None
         module_file = Path(module.__file__)
         module_dir = module_file.parent if module_file.is_file() else module_file
     except Exception as e:
-        logger.error(f"Failed to load module {args.module_name} for code generation: {e}")
+        logger.error(f"Failed to load module {module_name} for code generation: {e}")
         return
 
-    client_path = module_dir / "client.py"
+    client_path = module_dir / client_subdir / "client.py"
 
-    rpc_classes = get_custom_classes(args.module_name)
+    rpc_classes = get_custom_classes(module_name)
     logger.info(f"Obtained classes with RPC objects: {rpc_classes!r}")
 
-    generator = ClientGenerator(base=args.module_name == "bec_widgets")
-    logger.info(f"Generating client.py")
+    generator = ClientGenerator(base=module_name == "bec_widgets")
+    logger.info(f"Generating client file at {client_path}")
     generator.generate_client(rpc_classes)
     generator.write(str(client_path))
 
+    if module_name != "bec_widgets":
+        non_overwrite_classes = list(clsinfo.name for clsinfo in get_custom_classes("bec_widgets"))
+        logger.info(
+            f"Not writing plugins which would conflict with builtin classes: {non_overwrite_classes}"
+        )
+    else:
+        non_overwrite_classes = []
+
     for cls in rpc_classes.plugins:
-        logger.info(f"Writing plugins for: {cls}")
+        logger.info(f"Writing bec-designer plugin files for {cls.__name__}...")
+
+        if cls.__name__ in non_overwrite_classes:
+            logger.error(
+                f"Not writing plugin files for {cls.__name__} because a built-in widget with that name exists"
+            )
+
         plugin = DesignerPluginGenerator(cls)
         if not hasattr(plugin, "info"):
             continue
@@ -239,7 +302,9 @@ def main():
             return os.path.exists(os.path.join(plugin.info.base_path, file))
 
         if any(_exists(file) for file in plugin_filenames(plugin.info.plugin_name_snake)):
-            logger.debug(f"Skipping {plugin.info.plugin_name_snake} - a file already exists.")
+            logger.debug(
+                f"Skipping generation of extra plugin files for {plugin.info.plugin_name_snake} - at least one file out of 'plugin.py', 'pyproject', and 'register_{plugin.info.plugin_name_snake}.py' already exists."
+            )
             continue
 
         plugin.run()