bec_widgets/bec_widgets/cli/generate_cli.py

# pylint: disable=missing-module-docstring
from __future__ import annotations

import argparse
import inspect
import os
import sys

import black
import isort
from qtpy.QtCore import Property as QtProperty

from bec_widgets.utils.generate_designer_plugin import DesignerPluginGenerator
from bec_widgets.utils.plugin_utils import BECClassContainer, get_custom_classes

if sys.version_info >= (3, 11):
    from typing import get_overloads
else:
    print(
        "Python version is less than 3.11, using dummy function for get_overloads. "
        "If you want to use the real function 'typing.get_overloads()', please use Python 3.11 or later."
    )

    def get_overloads(_obj):
        """
        Dummy function for Python versions before 3.11.
        """
        return []


class ClientGenerator:
    def __init__(self):
        self.header = """# This file was automatically generated by generate_cli.py\n
from __future__ import annotations
import enum
from typing import Literal, Optional, overload

from bec_widgets.cli.rpc.rpc_base import RPCBase, rpc_call

# pylint: skip-file"""

        self.content = ""

    def generate_client(self, class_container: BECClassContainer):
        """
        Generate the client for the published classes.

        Args:
            class_container: The class container with the classes to generate the client for.
        """
        rpc_top_level_classes = class_container.rpc_top_level_classes
        rpc_top_level_classes.sort(key=lambda x: x.__name__)
        connector_classes = class_container.connector_classes
        connector_classes.sort(key=lambda x: x.__name__)

        self.write_client_enum(rpc_top_level_classes)
        for cls in connector_classes:
            self.content += "\n\n"
            self.generate_content_for_class(cls)

    def write_client_enum(self, published_classes: list[type]):
        """
        Write the client enum to the content.
        """
        self.content += """
class Widgets(str, enum.Enum):
    \"\"\"
    Enum for the available widgets.
    \"\"\"
    """
        for cls in published_classes:
            self.content += f'{cls.__name__} = "{cls.__name__}"\n    '

    def generate_content_for_class(self, cls):
        """
        Generate the content for the class.

        Args:
            cls: The class for which to generate the content.
        """

        class_name = cls.__name__

        # Generate the content
        if cls.__name__ == "BECDockArea":
            self.content += f"""
class {class_name}(RPCBase):"""
        else:
            self.content += f"""
class {class_name}(RPCBase):"""
        if not cls.USER_ACCESS:
            self.content += """...
    """
        for method in cls.USER_ACCESS:
            is_property_setter = False
            obj = getattr(cls, method, None)
            if obj is None:
                obj = getattr(cls, method.split(".setter")[0], None)
                is_property_setter = True
                method = method.split(".setter")[0]
            if obj is None:
                raise AttributeError(
                    f"Method {method} not found in class {cls.__name__}. Please check the USER_ACCESS list."
                )
            if isinstance(obj, (property, QtProperty)):
                # for the cli, we can map qt properties to regular properties
                if is_property_setter:
                    self.content += f"""
    @{method}.setter
    @rpc_call"""
                else:
                    self.content += """
    @property
    @rpc_call"""

                sig = str(inspect.signature(obj.fget))
                doc = inspect.getdoc(obj.fget)
            else:
                sig = str(inspect.signature(obj))
                doc = inspect.getdoc(obj)
                overloads = get_overloads(obj)
                for overload in overloads:
                    sig_overload = str(inspect.signature(overload))
                    self.content += f"""
    @overload
    def {method}{str(sig_overload)}: ...
    """

                self.content += """
    @rpc_call"""
            self.content += f"""
    def {method}{str(sig)}:
        \"\"\"
{doc}
        \"\"\""""

    def write(self, file_name: str):
        """
        Write the content to a file, automatically formatted with black.

        Args:
            file_name(str): The name of the file to write to.
        """
        # Combine header and content, then format with black
        full_content = self.header + "\n" + self.content
        try:
            formatted_content = black.format_str(full_content, mode=black.FileMode(line_length=100))
        except black.NothingChanged:
            formatted_content = full_content

        isort.Config(
            profile="black",
            line_length=100,
            multi_line_output=3,
            include_trailing_comma=True,
            known_first_party=["bec_widgets"],
        )
        formatted_content = isort.code(formatted_content)

        with open(file_name, "w", encoding="utf-8") as file:
            file.write(formatted_content)


def main():
    """
    Main entry point for the script, controlled by command line arguments.
    """

    parser = argparse.ArgumentParser(description="Auto-generate the client for RPC widgets")
    parser.add_argument("--core", action="store_true", help="Whether to generate the core client")

    args = parser.parse_args()

    if args.core:
        current_path = os.path.dirname(__file__)
        client_path = os.path.join(current_path, "client.py")

        rpc_classes = get_custom_classes("bec_widgets")

        generator = ClientGenerator()
        generator.generate_client(rpc_classes)
        generator.write(client_path)

        for cls in rpc_classes.plugins:
            plugin = DesignerPluginGenerator(cls)
            if not hasattr(plugin, "info"):
                continue

            # if the class directory already has a register, plugin and pyproject file, skip
            if os.path.exists(
                os.path.join(plugin.info.base_path, f"register_{plugin.info.plugin_name_snake}.py")
            ):
                continue
            if os.path.exists(
                os.path.join(plugin.info.base_path, f"{plugin.info.plugin_name_snake}_plugin.py")
            ):
                continue
            if os.path.exists(
                os.path.join(plugin.info.base_path, f"{plugin.info.plugin_name_snake}.pyproject")
            ):
                continue
            plugin.run()


if __name__ == "__main__":  # pragma: no cover
    sys.argv = ["generate_cli.py", "--core"]
    main()