"""Sources Controller Module""" __docformat__ = "numpy" # IMPORTATION STANDARD import argparse import json import logging import os from pathlib import Path from typing import Dict, List, Optional import openbb_terminal.core.session.hub_model as Hub import openbb_terminal.core.session.local_model as Local from openbb_terminal.core.session.constants import SOURCES_URL # IMPORTATION THIRDPARTY # IMPORTATION INTERNAL from openbb_terminal.core.session.current_user import ( get_current_user, is_local, set_sources, ) from openbb_terminal.core.session.hub_model import upload_config from openbb_terminal.custom_prompt_toolkit import NestedCompleter from openbb_terminal.decorators import log_start_end from openbb_terminal.menu import session from openbb_terminal.parent_classes import BaseController from openbb_terminal.rich_config import MenuText, console # pylint: disable=too-many-lines,no-member,too-many-public-methods,C0302 # pylint:disable=import-outside-toplevel logger = logging.getLogger(__name__) def unique(sequence): seen = set() return [x for x in sequence if not (x in seen or seen.add(x))] class SourcesController(BaseController): """Sources Controller class""" CHOICES_COMMANDS: List[str] = [ "get", "set", ] PATH = "/sources/" def __init__(self, queue: Optional[List[str]] = None): """Constructor""" super().__init__(queue) self.commands_with_sources: Dict[str, List[str]] = {} self.json_doc = {} if is_local(): self.load_sources_json_file() self.json_doc = get_current_user().sources.sources_dict self.generate_commands_with_sources() if session and get_current_user().preferences.USE_PROMPT_TOOLKIT: choices: dict = {c: {} for c in self.controller_choices} choices["get"] = {c: None for c in list(self.commands_with_sources.keys())} choices["set"] = {c: None for c in list(self.commands_with_sources.keys())} for cmd in list(self.commands_with_sources.keys()): choices["set"][cmd] = {c: None for c in self.commands_with_sources[cmd]} self.completer = NestedCompleter.from_nested_dict(choices) def load_sources_json_file(self): """Load the .json file""" SOURCES_FILE = Path(get_current_user().preferences.PREFERRED_DATA_SOURCE_FILE) if ( not os.getenv("TEST_MODE") and SOURCES_FILE.exists() and SOURCES_FILE.stat().st_size > 0 ): with open(str(SOURCES_FILE)) as json_file: self.json_doc = json.load(json_file) set_sources(self.json_doc) def generate_commands_with_sources(self): """Generate choices""" for context in self.json_doc: for menu in self.json_doc[context]: if isinstance(self.json_doc[context][menu], Dict): for submenu in self.json_doc[context][menu]: if isinstance(self.json_doc[context][menu][submenu], Dict): for subsubmenu in self.json_doc[context][menu][submenu]: self.commands_with_sources[ f"{context}_{menu}_{submenu}_{subsubmenu}" ] = self.json_doc[context][menu][submenu][subsubmenu] else: self.commands_with_sources[ f"{context}_{menu}_{submenu}" ] = self.json_doc[context][menu][submenu] else: self.commands_with_sources[f"{context}_{menu}"] = self.json_doc[ context ][menu] def print_help(self): """Print help""" mt = MenuText("sources/") mt.add_param( "_source", get_current_user().preferences.PREFERRED_DATA_SOURCE_FILE if is_local() else SOURCES_URL, ) mt.add_raw("\n") mt.add_info("_info_") mt.add_cmd("get") mt.add_cmd("set") console.print(text=mt.menu_text, menu="Data Sources") @log_start_end(log=logger) def call_get(self, other_args): """Process get command""" parser = argparse.ArgumentParser( add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter, prog="get", description="Get sources associated with a command and the one selected by default, using 'get '.", ) parser.add_argument( "-c", "--cmd", action="store", dest="cmd", choices=list(self.commands_with_sources.keys()), required="-h" not in other_args, help="Command that we want to check the available data sources and the default one.", metavar="COMMAND", ) if other_args and "-" not in other_args[0][0]: other_args.insert(0, "-c") ns_parser = self.parse_simple_args(parser, other_args) if ns_parser: if is_local(): self.load_sources_json_file() else: response = Hub.fetch_user_configs( get_current_user().profile.get_session() ) if response: configs = json.loads(response.content) Local.set_sources_from_hub(configs) self.json_doc = get_current_user().sources.sources_dict self.generate_commands_with_sources() try: the_item = self.commands_with_sources[ns_parser.cmd] except KeyError: console.print( [f"[red]'{ns_parser.cmd}' is not a valid command.[/red]\n"] ) return if the_item: console.print( f"\n[param]Default :[/param] {self.commands_with_sources[ns_parser.cmd][0]}" ) console.print( f"[param]Available :[/param] {', '.join(self.commands_with_sources[ns_parser.cmd])}\n" ) else: console.print("This command has no data sources available.\n") # pylint: disable=R0912 @log_start_end(log=logger) def call_set(self, other_args): """Process set command""" parser = argparse.ArgumentParser( add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter, prog="set", description="Set a default data sources associated with a command, using 'set '.", ) parser.add_argument( "-c", "--cmd", action="store", dest="cmd", choices=list(self.commands_with_sources.keys()), help="Command that we to select the default data source.", metavar="COMMAND", ) parser.add_argument( "-s", "--source", action="store", dest="source", type=str, help="Data source to use by default on specified command.", ) if other_args and "-" not in other_args[0][0]: other_args.insert(0, "-c") if "-s" not in other_args and "--source" not in other_args: other_args.insert(2, "-s") ns_parser = self.parse_simple_args(parser, other_args) if ns_parser: self.load_sources_json_file() success, valid_sources = self.update_dict_and_valid_sources(ns_parser) if success: set_sources(self.json_doc) if is_local(): try: self.write_json() self.update_dictionary() console.print( "[green]The data source was specified successfully.\n[/green]" ) except Exception as e: console.print( f"[red]Failed to write preferred data sources to file: " f"{get_current_user().preferences.PREFERRED_DATA_SOURCE_FILE}[/red]" ) console.print(f"[red]{e}[/red]") else: key = str(ns_parser.cmd).replace("_", "/") valid = self.commands_with_sources[ns_parser.cmd] valid.remove(ns_parser.source) valid.insert(0, ns_parser.source) value = ", ".join(valid) upload_config( key=key, value=value, type_="sources", auth_header=get_current_user().profile.get_auth_header(), ) else: console.print( f"[red]The data source selected is not valid, select one from: {', '.join(valid_sources)}.\n[/red]" ) def update_dict_and_valid_sources(self, ns_parser): """Update dictionary and valid sources""" menus = ns_parser.cmd.split("_") num_menus = len(menus) valid_sources = list() success = True # TODO: Change this to be recursive and handle any number of menus # Update dictionary if num_menus == 1: if ns_parser.source not in self.json_doc[menus[0]]: success = False valid_sources = self.json_doc[menus[0]] else: self.json_doc[menus[0]] = unique( [ns_parser.source] + self.json_doc[menus[0]] ) elif num_menus == 2: if ns_parser.source not in self.json_doc[menus[0]][menus[1]]: success = False valid_sources = self.json_doc[menus[0]][menus[1]] else: self.json_doc[menus[0]][menus[1]] = unique( [ns_parser.source] + self.json_doc[menus[0]][menus[1]] ) elif num_menus == 3: if ns_parser.source not in self.json_doc[menus[0]][menus[1]][menus[2]]: success = False valid_sources = self.json_doc[menus[0]][menus[1]][menus[2]] else: self.json_doc[menus[0]][menus[1]][menus[2]] = unique( [ns_parser.source] + self.json_doc[menus[0]][menus[1]][menus[2]] ) elif num_menus == 4: if ( ns_parser.source not in self.json_doc[menus[0]][menus[1]][menus[2]][menus[3]] ): success = False valid_sources = self.json_doc[menus[0]][menus[1]][menus[2]][menus[3]] else: self.json_doc[menus[0]][menus[1]][menus[2]][menus[3]] = unique( [ns_parser.source] + self.json_doc[menus[0]][menus[1]][menus[2]][menus[3]] ) return success, valid_sources def write_json(self): """Write json to file""" with open( get_current_user().preferences.PREFERRED_DATA_SOURCE_FILE, "w", ) as f: json.dump(self.json_doc, f, indent=4) def update_dictionary(self): """Update dictionary with new data sources""" for context in self.json_doc: for menu in self.json_doc[context]: if isinstance(self.json_doc[context][menu], Dict): for submenu in self.json_doc[context][menu]: if isinstance(self.json_doc[context][menu][submenu], Dict): for subsubmenu in self.json_doc[context][menu][submenu]: self.commands_with_sources[ f"{context}_{menu}_{submenu}_{subsubmenu}" ] = self.json_doc[context][menu][submenu][subsubmenu] else: self.commands_with_sources[ f"{context}_{menu}_{submenu}" ] = self.json_doc[context][menu][submenu] else: self.commands_with_sources[f"{context}_{menu}"] = self.json_doc[ context ][menu]