from argparse import SUPPRESS, ArgumentParser from contextlib import contextmanager from inspect import isfunction, unwrap from types import MethodType from typing import Callable, List from unittest.mock import patch from openbb_terminal.core.session.current_system import get_current_system from openbb_terminal.helper_funcs import check_file_type_saved, check_positive from openbb_terminal.rich_config import get_ordered_list_sources def __mock_parse_known_args_and_warn( controller, parser: ArgumentParser, other_args: List[str], export_allowed: int = 0, raw: bool = False, limit: int = 0, ) -> None: """Add the arguments that would have normally added by : - openbb_terminal.parent_classes.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: int Choose from 0, 1, 2 and EXPORT_BOTH_RAW_DATA_AND_FIGURES 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 > 0: choices_export = [] help_export = "Does not export!" if export_allowed == 1: choices_export = ["csv", "json", "xlsx"] help_export = "Export raw data into csv, json, xlsx" elif export_allowed == 2: 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, ) sources = get_ordered_list_sources(f"{controller.PATH}{parser.prog}") # Allow to change source if there is more than one if len(sources) > 1: parser.add_argument( "--source", action="store", dest="source", choices=sources, default=sources[0], # the first source from the list is the default help="Data source to select from", ) def __mock_parse_simple_args(parser: ArgumentParser, other_args: List[str]) -> None: """Add the arguments that would have normally added by: - openbb_terminal.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 Terminal 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 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 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_terminal.rich_config.ConsoleAndPanel.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 get_current_system().DEBUG_MODE: rich.start() patched_function_list = [] for patcher in patcher_list: patched_function_list.append(patcher.start()) yield patched_function_list if not get_current_system().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 Terminal 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 = patched_function.call_args.args argument_parser = 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: 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: command_list = controller.CHOICES_COMMANDS controller_choice_map: dict = {c: {} for c in controller.controller_choices} controller_choice_map["support"] = controller.SUPPORT_CHOICES controller_choice_map["about"] = controller.ABOUT_CHOICES controller_choice_map["hold"] = controller.HELP_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 get_current_system().DEBUG_MODE: raise Exception( f"On command : `{command}`.\n{str(exception)}" ) from exception return controller_choice_map