diff options
Diffstat (limited to 'cli/openbb_cli/controllers/choices.py')
| -rw-r--r-- | cli/openbb_cli/controllers/choices.py | 322 |
1 files changed, 322 insertions, 0 deletions
diff --git a/cli/openbb_cli/controllers/choices.py b/cli/openbb_cli/controllers/choices.py new file mode 100644 index 00000000000..7e66660ff79 --- /dev/null +++ b/cli/openbb_cli/controllers/choices.py @@ -0,0 +1,322 @@ +"""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 + + +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 |