openbb-cli 1.0.0__py3-none-any.whl

openbb_cli/controllers/base_platform_controller.py

@@ -0,0 +1,349 @@
+"""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 import obb
+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
+
+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: 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 not df.empty
+                    ):
+                        sheet_name = getattr(ns_parser, "sheet_name", None)
+                        if sheet_name and isinstance(sheet_name, list):
+                            sheet_name = sheet_name[0]
+
+                        export_data(
+                            export_type=",".join(ns_parser.export),
+                            dir_path=os.path.dirname(os.path.abspath(__file__)),
+                            func_name=translator.func.__name__,
+                            df=df,
+                            sheet_name=sheet_name,
+                            figure=fig,
+                        )
+                    elif hasattr(ns_parser, "export") and ns_parser.export and df.empty:
+                        session.console.print("[yellow]No data to export.[/yellow]")
+
+                except Exception as e:
+                    session.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(  # type: ignore
+            partial(bound_method, translator=translator), method
+        )
+        setattr(self, f"call_{name}", bound_method)
+
+    def _generate_controller_call(self, controller, name, parent_path, translators):
+        """Generate controller call."""
+
+        def method(self, _, controller, name, parent_path, translators):
+            """Call the controller."""
+            self.queue = self.load_class(
+                class_ins=controller,
+                name=name,
+                parent_path=parent_path,
+                translators=translators,
+                queue=self.queue,
+            )
+
+        # Bind the method to the class
+        bound_method = MethodType(method, self)
+
+        # Update the wrapper and set the attribute
+        bound_method = update_wrapper(  # type: ignore
+            partial(
+                bound_method,
+                name=name,
+                parent_path=parent_path,
+                translators=translators,
+                controller=controller,
+            ),
+            method,
+        )
+        setattr(self, f"call_{name}", bound_method)
+
+    def _get_command_description(self, command: str) -> str:
+        """Get command description."""
+        command_description = (
+            obb.reference["paths"]  # type: ignore
+            .get(f"{self.PATH}{command}", {})
+            .get("description", "")
+        )
+
+        if not command_description:
+            trl = self.translators.get(
+                f"{self._name}_{command}"
+            ) or self.translators.get(command)
+            if trl and hasattr(trl, "parser"):
+                command_description = trl.parser.description
+
+        return command_description.split(".")[0].lower()
+
+    def _get_menu_description(self, menu: str) -> str:
+        """Get menu description."""
+
+        def _get_sub_menu_commands():
+            """Get sub menu commands."""
+            sub_path = f"{self.PATH[1:].replace('/','_')}{menu}"
+            commands = []
+            for trl in self.translators:
+                if sub_path in trl:
+                    commands.append(trl.replace(f"{sub_path}_", ""))
+            return commands
+
+        menu_description = (
+            obb.reference["routers"]  # type: ignore
+            .get(f"{self.PATH}{menu}", {})
+            .get("description", "")
+        ) or ""
+        if menu_description:
+            return menu_description.split(".")[0].lower()
+
+        # If no description is found, return the sub menu commands
+        return ", ".join(_get_sub_menu_commands())
+
+    def print_help(self):
+        """Print help."""
+        mt = MenuText(self.PATH)
+
+        if self.CHOICES_MENUS:
+            for menu in self.CHOICES_MENUS:
+                description = self._get_menu_description(menu)
+                mt.add_menu(name=menu, description=description)
+
+            if self.CHOICES_COMMANDS:
+                mt.add_raw("\n")
+
+        if self.CHOICES_COMMANDS:
+            for command in self.CHOICES_COMMANDS:
+                command_description = self._get_command_description(command)
+                mt.add_cmd(
+                    name=command.replace(f"{self._name}_", ""),
+                    description=command_description,
+                )
+
+        if session.obbject_registry.obbjects:
+            mt.add_info("\nCached Results")
+            for key, value in list(session.obbject_registry.all.items())[
+                : session.settings.N_TO_DISPLAY_OBBJECT_REGISTRY
+            ]:
+                mt.add_raw(
+                    f"[yellow]OBB{key}[/yellow]: {value['command']}",
+                    left_spacing=True,
+                )
+
+        session.console.print(text=mt.menu_text, menu=self.PATH)
+
+        settings = session.settings
+        dev_mode = settings.DEBUG_MODE or settings.TEST_MODE
+        if mt.warnings and dev_mode:
+            session.console.print("")
+            for w in mt.warnings:
+                w_str = str(w).replace("{", "").replace("}", "").replace("'", "")
+                session.console.print(f"[yellow]{w_str}[/yellow]")
+            session.console.print("")

openbb_cli/controllers/choices.py

@@ -0,0 +1,324 @@
+"""This module contains functions to build the choice map for the controllers."""
+
+from argparse import SUPPRESS, ArgumentParser
+from contextlib import contextmanager
+from inspect import isfunction, unwrap
+from types import MethodType
+from typing import Callable, List, Literal
+from unittest.mock import patch
+
+from openbb_cli.controllers.utils import (
+    check_file_type_saved,
+    check_positive,
+)
+from openbb_cli.session import Session
+
+session = Session()
+
+
+def __mock_parse_known_args_and_warn(
+    controller,  # pylint: disable=unused-argument
+    parser: ArgumentParser,
+    other_args: List[str],
+    export_allowed: Literal[
+        "no_export", "raw_data_only", "figures_only", "raw_data_and_figures"
+    ] = "no_export",
+    raw: bool = False,
+    limit: int = 0,
+) -> None:
+    """Add arguments.
+
+    Add the arguments that would have normally added by :
+        - openbb_cli.base_controller.BaseController.parse_known_args_and_warn
+
+    Parameters
+    ----------
+    parser: argparse.ArgumentParser
+        Parser with predefined arguments
+    other_args: List[str]
+        list of arguments to parse
+    export_allowed: Literal["no_export", "raw_data_only", "figures_only", "raw_data_and_figures"]
+            Export options
+    raw: bool
+        Add the --raw flag
+    limit: int
+        Add a --limit flag with this number default
+    """
+    _ = other_args
+
+    parser.add_argument(
+        "-h", "--help", action="store_true", help="show this help message"
+    )
+
+    if export_allowed != "no_export":
+        choices_export = []
+        help_export = "Does not export!"
+
+        if export_allowed == "raw_data_only":
+            choices_export = ["csv", "json", "xlsx"]
+            help_export = "Export raw data into csv, json, xlsx"
+        elif export_allowed == "figures_only":
+            choices_export = ["png", "jpg", "pdf", "svg"]
+            help_export = "Export figure into png, jpg, pdf, svg "
+        else:
+            choices_export = ["csv", "json", "xlsx", "png", "jpg", "pdf", "svg"]
+            help_export = "Export raw data into csv, json, xlsx and figure into png, jpg, pdf, svg "
+
+        parser.add_argument(
+            "--export",
+            default="",
+            type=check_file_type_saved(choices_export),
+            dest="export",
+            help=help_export,
+            choices=choices_export,
+        )
+
+    if raw:
+        parser.add_argument(
+            "--raw",
+            dest="raw",
+            action="store_true",
+            default=False,
+            help="Flag to display raw data",
+        )
+    if limit > 0:
+        parser.add_argument(
+            "-l",
+            "--limit",
+            dest="limit",
+            default=limit,
+            help="Number of entries to show in data.",
+            type=check_positive,
+        )
+
+
+def __mock_parse_simple_args(parser: ArgumentParser, other_args: List[str]) -> None:
+    """Add arguments.
+
+    Add the arguments that would have normally added by:
+        - openbb_cli.parent_classes.BaseController.parse_simple_args
+
+    Parameters
+    ----------
+    parser: argparse.ArgumentParser
+        Parser with predefined arguments
+    other_args: List[str]
+        List of arguments to parse
+    """
+    parser.add_argument(
+        "-h", "--help", action="store_true", help="show this help message"
+    )
+    _ = other_args
+
+
+def __get_command_func(controller, command: str):
+    """Get the function with the name `f"call_{command}"` from controller object.
+
+    Parameters
+    ----------
+    controller: BaseController
+        Instance of the CLI Controller.
+    command: str
+        A name from controller.CHOICES_COMMANDS
+
+    Returns
+    -------
+    Callable: Command function.
+    """
+    if command not in controller.CHOICES_COMMANDS:
+        raise AttributeError(
+            f"The following command is not inside `CHOICES_COMMANDS` : '{command}'"
+        )
+
+    command = f"call_{command}"
+    command_func = getattr(controller, command)
+    command_func = unwrap(func=command_func)
+
+    if isfunction(command_func):
+        command_func = MethodType(command_func, controller)
+
+    return command_func
+
+
+def contains_functions_to_patch(command_func: Callable) -> bool:
+    """Check command function.
+
+    Check if a `command_func` actually contains the functions we want to mock, i.e.:
+        - parse_simple_args
+        - parse_known_args_and_warn
+
+    Parameters
+    ----------
+    command_func: Callable
+        Function to check.
+
+    Returns
+    -------
+    bool: Whether or not `command_func` contains the mocked functions.
+    """
+    co_names = command_func.__code__.co_names
+
+    return bool(
+        "parse_simple_args" in co_names or "parse_known_args_and_warn" in co_names
+    )
+
+
+@contextmanager
+def __patch_controller_functions(controller):
+    """Patch controller functions.
+
+    Patch the following function from a BaseController instance:
+        - parse_simple_args
+        - parse_known_args_and_warn
+
+    These functions take an 'argparse.ArgumentParser' object as parameter.
+    We want to intercept this 'argparse.ArgumentParser' object.
+
+    Parameters
+    ----------
+    controller: BaseController
+        BaseController object that needs to be patched.
+
+    Returns
+    -------
+    List[Callable]: List of mocked functions.
+    """
+    bound_mock_parse_known_args_and_warn = MethodType(
+        __mock_parse_known_args_and_warn,
+        controller,
+    )
+
+    rich = patch(
+        target="openbb_cli.config.console.Console.print",
+        return_value=None,
+    )
+
+    patcher_list = [
+        patch.object(
+            target=controller,
+            attribute="parse_simple_args",
+            side_effect=__mock_parse_simple_args,
+            return_value=None,
+        ),
+        patch.object(
+            target=controller,
+            attribute="parse_known_args_and_warn",
+            side_effect=bound_mock_parse_known_args_and_warn,
+            return_value=None,
+        ),
+    ]
+
+    if not session.settings.DEBUG_MODE:
+        rich.start()
+    patched_function_list = []
+    for patcher in patcher_list:
+        patched_function_list.append(patcher.start())
+
+    yield patched_function_list
+
+    if not session.settings.DEBUG_MODE:
+        rich.stop()
+    for patcher in patcher_list:
+        patcher.stop()
+
+
+def _get_argument_parser(
+    controller,
+    command: str,
+) -> ArgumentParser:
+    """Intercept the ArgumentParser instance from the command function.
+
+    A command function being a function starting with `call_`, like:
+        - call_help
+        - call_overview
+        - call_load
+
+    Parameters
+    ----------
+    controller: BaseController
+        Instance of the CLI Controller.
+    command: str
+        A name from `controller.CHOICES_COMMANDS`.
+
+    Returns
+    -------
+    ArgumentParser: ArgumentParser instance from the command function.
+    """
+    command_func: Callable = __get_command_func(controller=controller, command=command)
+
+    if not contains_functions_to_patch(command_func=command_func):
+        raise AssertionError(
+            f"One of these functions should be inside `call_{command}`:\n"
+            " - parse_simple_args\n"
+            " - parse_known_args_and_warn\n"
+        )
+
+    with __patch_controller_functions(controller=controller) as patched_function_list:
+        command_func([])
+
+        call_count = 0
+        for patched_function in patched_function_list:
+            call_count += patched_function.call_count
+            if patched_function.call_count == 1:
+                args, kwargs = patched_function.call_args
+                argument_parser = (
+                    kwargs["parser"] if kwargs.get("parser", None) else args[0]
+                )
+
+        if call_count != 1:
+            raise AssertionError(
+                f"One of these functions should be called once inside `call_{command}`:\n"
+                " - parse_simple_args\n"
+                " - parse_known_args_and_warn\n"
+            )
+
+    return argument_parser
+
+
+def _build_command_choice_map(argument_parser: ArgumentParser) -> dict:
+    """Build the choice map for a command."""
+    choice_map: dict = {}
+    for action in argument_parser._actions:  # pylint: disable=protected-access
+        if action.help == SUPPRESS:
+            continue
+        if len(action.option_strings) == 1:
+            long_name = action.option_strings[0]
+            short_name = ""
+        elif len(action.option_strings) == 2:
+            short_name = action.option_strings[0]
+            long_name = action.option_strings[1]
+        else:
+            raise AttributeError(f"Invalid argument_parser: {argument_parser}")
+
+        if hasattr(action, "choices") and action.choices:
+            choice_map[long_name] = {str(c): {} for c in action.choices}
+        else:
+            choice_map[long_name] = {}
+
+        if short_name and long_name:
+            choice_map[short_name] = long_name
+
+    return choice_map
+
+
+def build_controller_choice_map(controller) -> dict:
+    """Build the choice map for a controller."""
+    command_list = controller.CHOICES_COMMANDS
+    controller_choice_map: dict = {c: {} for c in controller.controller_choices}
+
+    for command in command_list:
+        try:
+            argument_parser = _get_argument_parser(
+                controller=controller,
+                command=command,
+            )
+            controller_choice_map[command] = _build_command_choice_map(
+                argument_parser=argument_parser
+            )
+        except Exception as exception:
+            if session.settings.DEBUG_MODE:
+                raise Exception(
+                    f"On command : `{command}`.\n{str(exception)}"
+                ) from exception
+
+    return controller_choice_map