| """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, Tuple |
| 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. |
| |
| 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 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. |
| |
| 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 |
| 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. |
| |
| 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, 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: |
| 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 |
|
|
