OpenBB/cli/openbb_cli/controllers/choices.py

"""Functions to build the choice map for the controllers."""

from argparse import SUPPRESS, ArgumentParser
from collections.abc import Callable
from contextlib import contextmanager
from inspect import isfunction, unwrap
from types import MethodType
from typing import Literal
from unittest.mock import patch

from openbb_cli.controllers.utils import (
    check_file_type_saved,
    check_positive,
    validate_register_key,
)
from openbb_cli.session import Session

session = Session()


def __mock_parse_known_args_and_warn(
    controller,
    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.

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

        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,
        )

    parser.add_argument(
        "--register_obbject",
        dest="register_obbject",
        action="store_false",
        default=True,
        help="Flag to store data in the OBBject registry, True by default.",
    )
    parser.add_argument(
        "--register_key",
        dest="register_key",
        default="",
        help="Key to reference data in the OBBject registry.",
        type=validate_register_key,
    )


def __mock_parse_simple_args(parser: ArgumentParser, other_args: list[str]) -> tuple:
    """Add arguments.

    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
    return None, None


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.

    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  # ty: ignore[unresolved-attribute]

    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.

    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, 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())

    try:
        yield patched_function_list
    finally:
        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.

    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:
        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