OpenBB/openbb_terminal/base_platform_controller.py

"""Platform Equity Controller."""

__docformat__ = "numpy"

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

import pandas as pd

from argparse_translator.argparse_class_processor import ArgparseClassProcessor
from openbb_terminal.core.session.current_user import get_current_user
from openbb_terminal.custom_prompt_toolkit import NestedCompleter
from openbb_terminal.helper_funcs import print_rich_table
from openbb_terminal.menu import session
from openbb_terminal.parent_classes import BaseController
from openbb_terminal.rich_config import MenuText, console

logger = logging.getLogger(__name__)


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,
        platform_target: Optional[type] = None,
        queue: Optional[List[str]] = None,
        translators: Optional[Dict] = None,
    ):
        """Construct a Platform based Controller."""
        self.PATH = 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)
            if platform_target
            else DummyTranslation()
        )
        self.translators = translators or self._translated_target.translators
        self.paths = self._translated_target.paths

        if self.translators:
            self._generate_commands()
            self._generate_sub_controllers()

            if session and get_current_user().preferences.USE_PROMPT_TOOLKIT:
                choices: dict = self.choices_default
                self.completer = NestedCompleter.from_nested_dict(choices)

    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,
                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):
        def method(self, other_args: List[str], translator=translator):
            parser = translator.parser

            if ns_parser := self.parse_known_args_and_warn(parser, other_args):
                try:
                    obbject = translator.execute_func(parsed_args=ns_parser)

                    if hasattr(ns_parser, "chart") and ns_parser.chart:
                        obbject.show()
                    elif hasattr(obbject, "to_dataframe"):
                        print_rich_table(obbject.to_dataframe())
                    elif isinstance(obbject, dict):
                        print_rich_table(
                            pd.DataFrame.from_dict(obbject, orient="index")
                        )
                    else:
                        console.print(obbject)

                except Exception as e:
                    console.print(f"[red]{e}[/]\n")
                    return

        # Bind the method to the class
        bound_method = MethodType(method, self)

        # Update the wrapper and set the attribute
        bound_method = update_wrapper(
            partial(bound_method, translator=translator), method
        )
        setattr(self, f"call_{name}", bound_method)

    def _generate_controller_call(self, controller, name, translators):
        def method(self, _, controller=controller, name=name, translators=translators):
            self.queue = self.load_class(
                controller, name, self.queue, translators=translators
            )

        # Bind the method to the class
        bound_method = MethodType(method, self)

        # Update the wrapper and set the attribute
        bound_method = update_wrapper(
            partial(
                bound_method, controller=controller, name=name, translators=translators
            ),
            method,
        )
        setattr(self, f"call_{name}", bound_method)

    def print_help(self):
        """Print help."""
        mt = MenuText(self._name, 80)

        if self.CHOICES_MENUS:
            mt.add_raw("Menus\n\n")
            for menu in self.CHOICES_MENUS:
                mt.add_menu(menu)

        if self.CHOICES_COMMANDS:
            mt.add_raw("\nCommands\n\n")
            for command in self.CHOICES_COMMANDS:
                mt.add_cmd(command.replace(f"{self._name}_", ""))

        console.print(text=mt.menu_text, menu=self._name)