path: root/cli/openbb_cli/controllers/base_platform_controller.py

"""Platform Equity Controller."""

import os
from functools import partial, update_wrapper
from types import MethodType
from typing import Dict, List, Optional

import pandas as pd
from openbb_charting.core.openbb_figure import OpenBBFigure
from openbb_cli.argparse_translator.argparse_class_processor import (
    ArgparseClassProcessor,
)
from openbb_cli.config.menu_text import MenuText
from openbb_cli.controllers.base_controller import BaseController
from openbb_cli.controllers.utils import export_data, print_rich_table
from openbb_cli.session import Session
from openbb_core.app.model.obbject import OBBject

from openbb import obb

session = Session()


class DummyTranslation:
    """Dummy Translation for testing."""

    def __init__(self):
        """Construct a Dummy Translation Class."""
        self.paths = {}
        self.translators = {}


class PlatformController(BaseController):
    """Platform Controller Base class."""

    CHOICES_GENERATION = True

    def __init__(
        self,
        name: str,
        parent_path: List[str],
        platform_target: Optional[type] = None,
        queue: Optional[List[str]] = None,
        translators: Optional[Dict] = None,
    ):
        """Construct a Platform based Controller."""
        self.PATH = f"/{'/'.join(parent_path)}/{name}/" if parent_path else f"/{name}/"
        super().__init__(queue)
        self._name = name

        if not (platform_target or translators):
            raise ValueError("Either platform_target or translators must be provided.")

        self._translated_target = (
            ArgparseClassProcessor(
                target_class=platform_target, reference=obb.reference["paths"]  # type: ignore
            )
            if platform_target
            else DummyTranslation()
        )
        self.translators = (
            translators
            if translators is not None
            else getattr(self._translated_target, "translators", {})
        )
        self.paths = getattr(self._translated_target, "paths", {})

        if self.translators:
            self._link_obbject_to_data_processing_commands()
            self._generate_commands()
            self._generate_sub_controllers()
            self.update_completer(self.choices_default)

    def _link_obbject_to_data_processing_commands(self):
        """Link data processing commands to OBBject registry."""
        for _, trl in self.translators.items():
            for action in trl._parser._actions:  # pylint: disable=protected-access
                if action.dest == "data":
                    action.choices = [
                        "OBB" + str(i)
                        for i in range(len(session.obbject_registry.obbjects))
                    ]
                    action.type = str
                    action.nargs = None

    def _intersect_data_processing_commands(self, ns_parser):
        """Intersect data processing commands and change the obbject id into an actual obbject."""
        if hasattr(ns_parser, "data"):
            ns_parser.data = int(ns_parser.data.replace("OBB", ""))
            if ns_parser.data in range(len(session.obbject_registry.obbjects)):
                obbject = session.obbject_registry.get(ns_parser.data)
                setattr(ns_parser, "data", obbject.results)

        return ns_parser

    def _generate_sub_controllers(self):
        """Handle paths."""
        for path, value in self.paths.items():
            if value == "path":
                continue

            sub_menu_translators = {}
            choices_commands = []

            for translator_name, translator in self.translators.items():
                if f"{self._name}_{path}" in translator_name:
                    new_name = translator_name.replace(f"{self._name}_{path}_", "")
                    sub_menu_translators[new_name] = translator
                    choices_commands.append(new_name)

                    if translator_name in self.CHOICES_COMMANDS:
                        self.CHOICES_COMMANDS.remove(translator_name)

            # Create the sub controller as a new class
            class_name = f"{self._name.capitalize()}{path.capitalize()}Controller"
            SubController = type(
                class_name,
                (PlatformController,),
                {
                    "CHOICES_GENERATION": True,
                    # "CHOICES_MENUS": [],
                    "CHOICES_COMMANDS": choices_commands,
                },
            )

            self._generate_controller_call(
                controller=SubController,
                name=path,
                parent_path=self.path,
                translators=sub_menu_translators,
            )

    def _generate_commands(self):
        """Generate commands."""
        for name, translator in self.translators.items():
            # Prepare the translator name to create a command call in the controller
            new_name = name.replace(f"{self._name}_", "")

            self._generate_command_call(name=new_name, translator=translator)

    def _generate_command_call(self, name, translator):
        """Generate command call."""

        def method(self, other_args: List[str], translator=translator):
            """Call the translator."""
            parser = translator.parser

            if ns_parser := self.parse_known_args_and_warn(
                parser=parser,
                other_args=other_args,
                export_allowed="raw_data_and_figures",
            ):
                try:
                    ns_parser = self._intersect_data_processing_commands(ns_parser)

                    obbject = translator.execute_func(parsed_args=ns_parser)
                    df: pd.DataFrame = pd.DataFrame()
                    fig: Optional[OpenBBFigure] = None
                    title = f"{self.PATH}{translator.func.__name__}"

                    if obbject:
                        if isinstance(obbject, OBBject):
                            if session.max_obbjects_exceeded() and obbject.results:
                                session.obbject_registry.remove()
                                session.console.print(
                                    "[yellow]Maximum number of OBBjects reached. The oldest entry was removed.[yellow]"
                                )

                            # use the obbject to store the command so we can display it later on results
                            obbject.extra["command"] = f"{title} {' '.join(other_args)}"

                            register_result = session.obbject_registry.register(obbject)

                            # we need to force to re-link so that the new obbject
                            # is immediately available for data processing commands
                            self._link_obbject_to_data_processing_commands()
                            # also update the completer
                            self.update_completer(self.choices_default)

                            if (
                                session.settings.SHOW_MSG_OBBJECT_REGISTRY
                                and register_result
                            ):
                                session.console.print(
                                    "Added `OBBject` to cached results."
                                )

                            # making the dataframe available
                            # either for printing or exporting (or both)
                            df = obbject.to_dataframe()

                            if hasattr(ns_parser, "chart") and ns_parser.chart:
                                obbject.show()
                                fig = obbject.chart.fig if obbject.chart else None
                            else:
                                if isinstance(df.columns, pd.RangeIndex):
                                    df.columns = [str(i) for i in df.columns]

                                print_rich_table(df=df, show_index=True, title=title)

                        elif isinstance(obbject, dict):
                            df = pd.DataFrame.from_dict(obbject, orient="columns")
                            print_rich_table(df=df, show_index=True, title=title)

                        elif not isinstance(obbject, OBBject):
                            session.console.print(obbject)

                    if (
                        hasattr(ns_parser, "export")
                        and ns_parser.export
                        and