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

"""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", "svg"]
            help_export = "Export figure into png, jpg, svg "
        else:
            choices_export = ["csv", "json", "xlsx", "png", "jpg", "svg"]
            help_export = (
                "Export raw data into csv, json, xlsx and figure into png, jpg, 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: