generate_sdk.py

import os
import re
import subprocess
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, TextIO

from openbb_terminal.rich_config import console
from openbb_terminal.sdk_core.sdk_helpers import get_sdk_imports_text
from openbb_terminal.sdk_core.trailmap import Trailmap, get_trailmaps

REPO_ROOT = Path(__file__).parent.joinpath("openbb_terminal").resolve()

sub_names = {
    "defi": "DeFi",
    "disc": "Discovery",
    "dd": "Due Diligence",
    "onchain": "OnChain",
    "ov": "Overview",
    "keys": "Keys",
    "forecast": "Forecasting",
    "alt": "Alternative",
    "crypto": "Cryptocurrency",
    "ba": "Behavioral Analysis",
    "ca": "Comparison Analysis",
    "dps": "Darkpool Shorts",
    "po": "Portfolio Optimization",
    "qa": "Quantitative Analysis",
    "screener": "Screener",
    "sia": "Sector Industry Analysis",
    "ta": "Technical Analysis",
    "th": "Trading Hours",
    "fa": "Fundamental Analysis",
    "Oss": "Open Source Software",
    "funds": "Mutual Funds",
    "gov": "Government",
    "ins": "Insiders",
    "nft": "NFT",
}


sdk_openbb_var = """
class SDKLogger:
    def __init__(self) -> None:
        self.__check_initialize_logging()

    def __check_initialize_logging(self):
        if not cfg.LOGGING_SUPPRESS:
            self.__initialize_logging()

    @staticmethod
    def __initialize_logging() -> None:
        # pylint: disable=C0415
        from openbb_terminal.core.log.generation.settings_logger import log_all_settings
        from openbb_terminal.loggers import setup_logging

        cfg.LOGGING_SUB_APP = "sdk"
        setup_logging()
        log_all_settings()


openbb = OpenBBSDK()\r\r
"""

disable_lines = "# flake8: noqa\r# pylint: disable=C0301,R0902,R0903\r"


class BuildCategoryModelClasses:
    def __init__(self, trailmaps: List[Trailmap]) -> None:
        self.trailmaps = trailmaps
        self.categories: Dict[str, Any] = {}
        self.import_modules: Dict[str, Any] = {}
        self.root_modules: Dict[str, Any] = {}
        self.import_cat_class = (
            "from openbb_terminal.sdk_core.sdk_helpers import Category\r"
        )

        for tmap in self.trailmaps:
            local = tmap.location_path
            self.categories = self.add_todict(self.categories, local, tmap)

        for folder in ["models", "controllers"]:
            if not (REPO_ROOT / f"sdk_core/{folder}").exists():
                (REPO_ROOT / f"sdk_core/{folder}").mkdir()
            if not (REPO_ROOT / f"sdk_core/{folder}/__init__.py").exists():
                with open(REPO_ROOT / f"sdk_core/{folder}/__init__.py", "w") as f:
                    f.write(
                        "# flake8: noqa\r# pylint: disable=unused-import,wrong-import-order\r\r"
                    )

        if not (REPO_ROOT / "sdk_core/__init__.py").exists():
            with open(REPO_ROOT / "sdk_core/__init__.py", "w") as f:
                f.write("")

    def add_todict(self, d: dict, location_path: list, tmap: Trailmap) -> dict:
        """Adds the trailmap to the dictionary. A trailmap is a path to a function
        in the sdk. This function creates the dictionary paths to the function."""

        if location_path[0] not in d:
            d[location_path[0]] = {}

        if len(location_path) > 1:
            self.add_todict(d[location_path[0]], location_path[1:], tmap)
        else:
            d[location_path[0]][tmap.class_attr] = tmap

        return d

    def get_nested_dict(self, d: dict) -> dict:
        """Gets the nested dictionary of the category."""
        nested_dict = {}
        for k, v in d.items():
            if isinstance(v, dict):
                nested_dict[k] = self.get_nested_dict(v)
        return nested_dict

    def get_subcat_fullname(self, cat: str) -> str:
        """Gets the full category name from the shortened category name."""
        if cat in sub_names:
            return sub_names[cat]
        return cat.title().replace(" ", "")

    def check_submodules(self, category: str) -> bool:
        """Checks if a category has submodules."""
        return any(isinstance(d, dict) for d in self.categories[category].values())

    def write_init_imports(self, importstr: str, filestr: str) -> None:
        """Checks if a category has submodules and adds the imports to the init file."""
        regex = re.compile(importstr)
        with open(REPO_ROOT / filestr) as init_file:
            check_init = bool(regex.search(init_file.read()))
            if not check_init:
                with open(REPO_ROOT / filestr, "a") as init_file:
                    init_file.write(f"{importstr}\r")

    def write_sdk_class(
        self, category: str, f: TextIO, cls_type: Optional[str] = ""
    ) -> None:
        """Writes the sdk class to the file.

        Parameters
        ----------
        category : str
            The category name.
        f : TextIO
            The file to write to.
        cls_type : Optional[str], optional
            The class type, by default ""
        """
        if category in self.root_modules:
            class_name = f"{category.title()}{cls_type}(model.{category.title()}Root)"
        else:
            class_name = f"{category.title()}{cls_type}"
        f.write(
            f'class {class_name}:\r    """{self.get_subcat_fullname(category)} Module.\r'
        )

    def write_class_property(
        self, category: str, f: TextIO, subcat: Optional[str] = ""
    ) -> None:
        """Writes the class property to the file.

        Parameters
        ----------
        category : str
            The category name.
        f : TextIO
            The file to write to.
        subcat : Optional[str], optional
            The subcategory name, by default None
        """
        def_name = subcat if subcat else category
        if subcat:
            subcat = f" {self.get_subcat_fullname(subcat)}"

        category = sub_names.get(category, category.title())

        f.write(
            f"    @property\r    def {def_name}(self):\r        "
            f'"""{category}{subcat} Submodule\r'
        )

    def write_class_attr_docs(
        self,
        d: dict,
        f: TextIO,
        module: bool = True,
        sdk_root: bool = False,
        trail: list = None,
    ) -> None:
        """Writes the class attribute docs to the category file.

        Parameters
        ----------
        d : dict
            The dictionary of the category.
        f : TextIO
            The file to write to.
        module : bool, optional
            If the category is a module, by default True
        sdk_root : bool, optional
            If this is the sdk root file, by default False
        trail : list, optional
            The trail to the function, by default None
        """
        add_indent = "" if module else "    "
        added_attributes = False

        for v in d.values():
            if isinstance(v, Trailmap):
                if not added_attributes:
                    f.write(f"\r{add_indent}    Attributes:\r")
                    added_attributes = True

                for key in ["model", "view"]:
                    if v.func_attrs.get(key, None) and v.func_attrs[key].short_doc:
                        view = v.view_name if key == "view" else ""
                        f.write(
                            f"{add_indent}        `{v.class_attr}{view}`: {v.func_attrs[key].short_doc}\\n\r"
                        )

        f.write(f'{add_indent}    """\r\r')
        if trail:
            f.write(f'    _location_path = "{".".join(trail)}"\r')
        if module:
            f.write("    def __init__(self):\r        super().__init__()\r")
        elif sdk_root:
            f.write("    __version__ = obbff.VERSION\r\r")
            f.write("    def __init__(self):\r        SDKLogger()\r")

    def write_class_attributes(
        self, d: dict, f: TextIO, cat: Optional[str] = None
    ) -> None:
        """Writes the class attributes to the category file.

        Parameters
        ----------
        d : dict
            The dictionary of the category.
        f : TextIO
            The file to write to.
        cat : Optional[str], optional
            The category name, by default None
        """
        add_indent = ""

        missing_deps = {
            "forecast": "FORECASTING_TOOLKIT",
            "po": "OPTIMIZATION_TOOLKIT",
        }

        if cat in missing_deps:
            f.write(
                f"""
        if not lib.{missing_deps[cat]}_ENABLED:\r
            # pylint: disable=C0415
            from openbb_terminal.rich_config import console
            console.print(lib.{missing_deps[cat]}_WARNING)\r\r"""
            )

            add_indent = "    "
            f.write(f"        if lib.{missing_deps[cat]}_ENABLED:\r")

        for v in d.values():
            if isinstance(v, Trailmap):
                for attr, func in zip(["", v.view_name], [v.model_func, v.view_func]):
                    if func:
                        f.write(
                            f"{add_indent}        self.{v.class_attr}{attr} = {func}\r"
                        )
        f.write("\r\r")

    def write_category(self, category: str, d: dict, f: TextIO) -> None:
        """Writes the category class to the file

        Parameters
        ----------
        category : str
            The category name
        d : dict
            The category dictionary
        f : TextIO
            The file to write to
        """
        subname = self.get_subcat_fullname(category)

        # If this catergory has no attributes, then we don't write it to the file.
        if not any(isinstance(v, Trailmap) for v in d.values()):
            return

        self.root_modules[category] = f"{category.title().replace(' ', '')}Root"

        f.write(f"class {self.root_modules[category]}(Category):\r")
        f.write(f'    """{subname.title()} Module\r')

        self.write_class_attr_docs(d, f, trail=[category])
        self.write_class_attributes(d, f, category)

    def write_nested_category(self, category: str, d: dict, f: TextIO) -> None:
        """Writes the nested category classes.

        Parameters
        ----------
        category : str
            The category name
        d : dict
            The category dictionary
        f : TextIO
            The file to write to
        """
        for nested_category, nested_dict in d.items():
            if isinstance(nested_dict, Trailmap):
                continue
            subname = self.get_subcat_fullname(nested_category)

            class_name = f"{category.title()}{subname.replace(' ', '')}(Category)"
            f.write(f'class {class_name}:\r    """{subname} Module.\r')

            self.write_nested_submodule_docs(nested_dict, f)

            if isinstance(nested_dict, dict):
                self.write_class_attr_docs(
                    nested_dict, f, trail=[category, nested_category]
                )
                self.write_class_attributes(nested_dict, f, nested_category)

    def write_submodule_doc(
        self, k: str, f: TextIO, added: bool = False, indent: bool = False
    ) -> None:
        """Writes the submodules to the class docstring.

        Parameters
        ----------
        d : dict
            The category dictionary
        f : TextIO
            The file to write to
        added : bool, optional
            Whether or not "Submodules:" have been added to the docstring, by default False
        indent : bool, optional
            Whether or not to add an indent to the docstring, by default False
        """
        add_indent = "    " if indent else ""
        if not added:
            f.write("\r        Submodules:\r")
        subcat_name = self.get_subcat_fullname(k)
        f.write(f"{add_indent}        `{k}`: {subcat_name} Module\r")

    def write_nested_submodule_docs(
        self, nested_dict: dict, f: TextIO, indent: bool = False
    ) -> None:
        """Writes the nested submodule docs to the class docstring.

        Parameters
        ----------
        nested_dict : dict
            The nested dictionary
        f : TextIO
            The file to write to
        indent : bool, optional
            Whether or not to add an indent to the docstring, by default False
        """
        added = False
        nested_subcat = self.get_nested_dict(nested_dict)
        if nested_subcat:
            for k in nested_subcat:
                if isinstance(nested_dict[k], Trailmap):
                    continue
                self.write_submodule_doc(k, f, added, indent)
                if not added:
                    added = True

    def write_category_file(self, category: str, d: dict) -> None:
        """Writes the category file. This is the file that contains the categories
        and subcategories of the sdk.

        Parameters
        ----------
        category : str
            The category name
        d : dict
            The category dictionary
        """
        with open(REPO_ROOT / f"sdk_core/models/{category}_sdk_model.py", "w") as f:
            import_cat_class = self.import_cat_class
            if category in self.root_modules:
                import_cat_class = ""
                category = self.root_modules[category]

            f.write(
                f"{disable_lines}{import_cat_class}import openbb_terminal.sdk_core.sdk_init as lib\r\r\r"
            )
            if category not in self.root_modules and any(
                isinstance(v, Trailmap) for v in d.values()
            ):
                self.write_init_imports(
                    f"from .{category}_sdk_model import {category.title()}Root",
                    "sdk_core/models/__init__.py",
                )

            self.write_category(category, d, f)
            self.write_nested_category(category, d, f)

            f.seek(f.tell() - 2, os.SEEK_SET)
            f.truncate()

    def write_sdk_controller_file(self, category: str, d: dict) -> None:
        """Writes the sdk controller file. This is the file that contains the
        controller classes for the sdk.

        Parameters
        ----------
        category : str
            The category name
        d : dict
            The category dictionary
        """
        added_init_imports = []
        with open(
            REPO_ROOT / f"sdk_core/controllers/{category}_sdk_controller.py", "w"
        ) as f:
            f.write(
                f"{disable_lines}from openbb_terminal.sdk_core.models "
                f"import {category}_sdk_model as model\r\r\r"
            )

            if category not in added_init_imports and any(
                isinstance(v, dict) for v in d.values()
            ):
                self.write_init_imports(
                    f"from .{category}_sdk_controller import {category.title()}Controller",
                    "sdk_core/controllers/__init__.py",
                )
                added_init_imports.append(category)

            self.write_sdk_class(category, f, "Controller")
            self.write_nested_submodule_docs(self.categories[category], f, True)
            self.write_class_attr_docs(d, f, False)

            for subcategory in self.categories[category]:
                if isinstance(d[subcategory], Trailmap):
                    continue

                self.write_class_property(category, f, subcategory)
                self.write_nested_submodule_docs(d[subcategory], f, True)
                self.write_class_attr_docs(d[subcategory], f, False)
                f.write(
                    f"        return model.{category.title()}"
                    f"{self.get_subcat_fullname(subcategory).replace(' ', '')}()\r\r"
                )

            f.seek(f.tell() - 1, os.SEEK_SET)
            f.truncate()

    def write_sdk_file(self) -> None:
        """Writes the main sdk file. This is the file that we initialize the SDK with openbb."""

        with open(REPO_ROOT / "sdk.py", "w") as f:
            f.write(
                f'{get_sdk_imports_text()}class OpenBBSDK:\r    """OpenBB SDK Class.\r'
            )
            root_attrs = self.categories.pop("root")
            if root_attrs:
                self.write_class_attr_docs(root_attrs, f, False, True)
                self.write_class_attributes(root_attrs, f)

            for category in self.categories:
                self.write_class_property(category, f)
                self.write_nested_submodule_docs(self.categories[category], f, True)
                self.write_class_attr_docs(self.categories[category], f, False)

                if self.check_submodules(category):
                    f.write(f"        return ctrl.{category.title()}Controller()\r\r")
                else:
                    f.write(f"        return model.{self.root_modules[category]}()\r\r")

            f.write("\r".join(sdk_openbb_var.splitlines()))

    def build(self) -> None:
        """Builds the SDK."""
        for path in ["sdk_core/models", "sdk_core/controllers"]:
            for file in (REPO_ROOT / path).glob("*.py"):
                if file.name == "__init__.py":
                    with open(file, "w") as f:
                        f.write(
                            "# flake8: noqa\r"
                            "# pylint: disable=unused-import,wrong-import-order\r\r"
                        )
                    continue
                file.unlink()

        for category, d in self.categories.items():
            if isinstance(d, Trailmap) or category == "root":
                continue
            self.write_category_file(category, d)

            if self.check_submodules(category):
                self.write_sdk_controller_file(category, d)
        self.write_sdk_file()

        for path in ["sdk.py", "sdk_core", "sdk_core/models", "sdk_core/controllers"]:
            for file in (REPO_ROOT / path).glob("*.py"):
                with open(file, "rb") as f:
                    content = f.read()
                with open(file, "wb") as f:
                    f.write(content.replace(b"\r", b"\n"))

        # We run black to make sure the code is formatted correctly
        subprocess.check_call(["black", "openbb_terminal"])


def generate_sdk(sort: bool = False) -> None:
    """Main function to generate the SDK.

    Parameters
    ----------
    sort : bool, optional
        Whether to sort the CSVs, by default False
    """
    trailmaps = get_trailmaps(sort)

    console.print("[yellow]Generating SDK...[/]")
    BuildCategoryModelClasses(trailmaps).build()
    console.print("[green]SDK Generated Successfully.[/]")
    return


if __name__ == "__main__":
    sort_csv = False
    if len(sys.argv) > 1:
        if sys.argv[1] == "sort":
            sort_csv = True
            console.print("\n\n[bright_magenta]Sorting CSV...[/]\n")
        else:
            console.print("[red]Invalid argument.\n Accepted arguments: sort[/]")

    generate_sdk(sort_csv)