From cdae8f49d2c3d3d2e7c0e445c715c736ca591002 Mon Sep 17 00:00:00 2001 From: Jeroen Bouma Date: Mon, 14 Mar 2022 13:47:00 +0100 Subject: [PATCH] Add multi-plotting functionality as well as robustness of functions --- gamestonk_terminal/economy/econdb_model.py | 52 +++ gamestonk_terminal/economy/econdb_view.py | 98 +++-- .../economy/economy_controller.py | 374 +++++++++++++++--- gamestonk_terminal/economy/fred_model.py | 41 +- gamestonk_terminal/economy/fred_view.py | 99 +++-- gamestonk_terminal/economy/plot_view.py | 69 ++++ gamestonk_terminal/economy/yfinance_view.py | 42 +- 7 files changed, 591 insertions(+), 184 deletions(-) create mode 100644 gamestonk_terminal/economy/plot_view.py diff --git a/gamestonk_terminal/economy/econdb_model.py b/gamestonk_terminal/economy/econdb_model.py index ed49f05b10a..c88dbe4b658 100644 --- a/gamestonk_terminal/economy/econdb_model.py +++ b/gamestonk_terminal/economy/econdb_model.py @@ -340,6 +340,58 @@ def get_macro_data( return df +@log_start_end(log=logger) +def get_aggregated_macro_data( + parameters: list, + countries: list, + start_date: int = None, + end_date: int = None, + convert_currency: str = "USD", +) -> pd.DataFrame: + """This functions groups the data queried from the EconDB database [Source: EconDB] + + Parameters + ---------- + parameters: list + The type of data you wish to acquire + countries : list + the selected country or countries + start_date : str + The starting date, format "YEAR-MONTH-DAY", i.e. 2010-12-31. + end_date : str + The end date, format "YEAR-MONTH-DAY", i.e. 2020-06-05. + convert_currency : str + In what currency you wish to convert all values. + + Returns + ---------- + pd.DataFrame + A DataFrame with the requested macro data of all chosen countries + """ + country_data: Dict[Any, Dict[Any, pd.Series]] = {} + + for country in countries: + country_data[country] = {} + for parameter in parameters: + country_data[country][parameter] = get_macro_data( + parameter, country, convert_currency + ) + + if country_data[country][parameter] is not None and start_date or end_date: + country_data[country][parameter] = country_data[country][parameter].loc[ + start_date:end_date + ] + + country_data_df = ( + pd.DataFrame.from_dict(country_data, orient="index").stack().to_frame() + ) + country_data_df = pd.DataFrame( + country_data_df[0].values.tolist(), index=country_data_df.index + ).T + + return country_data_df + + @log_start_end(log=logger) def get_treasuries( instruments: list, diff --git a/gamestonk_terminal/economy/econdb_view.py b/gamestonk_terminal/economy/econdb_view.py index a8365d84175..b5ac56c1210 100644 --- a/gamestonk_terminal/economy/econdb_view.py +++ b/gamestonk_terminal/economy/econdb_view.py @@ -1,9 +1,10 @@ """ EconDB View """ __docformat__ = "numpy" +# pylint:disable=too-many-arguments import logging import os -from typing import Optional, List, Dict, Any +from typing import Optional, List, Dict import pandas as pd from matplotlib import pyplot as plt @@ -58,41 +59,28 @@ def show_macro_data( ---------- Plots the Series. """ - country_data: Dict[Any, Dict[Any, pd.Series]] = {} + country_data_df = econdb_model.get_aggregated_macro_data( + parameters, countries, start_date, end_date, convert_currency + ) if external_axes is None: _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI) else: ax = external_axes[0] - for country in countries: - country_data[country] = {} - for parameter in parameters: - country_data[country][parameter] = econdb_model.get_macro_data( - parameter, country, convert_currency - ) - - if country_data[country][parameter] is not None and start_date or end_date: - country_data[country][parameter] = country_data[country][parameter].loc[ - start_date:end_date - ] - - df = pd.DataFrame.from_dict(country_data, orient="index").stack().to_frame() - df = pd.DataFrame(df[0].values.tolist(), index=df.index).T - - maximum_value = df.max().max() + maximum_value = country_data_df.max().max() if maximum_value > 1_000_000_000: - df_rounded = df / 1_000_000_000 + df_rounded = country_data_df / 1_000_000_000 denomination = f"[{convert_currency} Billions]" elif maximum_value > 1_000_000: - df_rounded = df / 1_000_000 + df_rounded = country_data_df / 1_000_000 denomination = f"[{convert_currency} Millions]" elif maximum_value > 1_000: - df_rounded = df / 1_000 + df_rounded = country_data_df / 1_000 denomination = f"[{convert_currency} Thousands]" else: - df_rounded = df + df_rounded = country_data_df denomination = f"[{convert_currency}]" for column in df_rounded.columns: @@ -114,6 +102,8 @@ def show_macro_data( ax.plot(df_rounded[column]) ax.set_title(f"{parameter_label} of {country_label} {denomination}") + df_rounded.columns = ["_".join(column) for column in df_rounded.columns] + if raw: print_rich_table( df_rounded.fillna("-").iloc[-10:], @@ -138,6 +128,7 @@ def show_treasuries( frequency: str, start_date: str = None, end_date: str = None, + store: bool = False, raw: bool = False, external_axes: Optional[List[plt.axes]] = None, export: str = "", @@ -156,6 +147,8 @@ def show_treasuries( The starting date, format "YEAR-MONTH-DAY", i.e. 2010-12-31. end_date : str The end date, format "YEAR-MONTH-DAY", i.e. 2020-06-05. + store : bool + Whether to prevent plotting the data. raw : bool Whether to display the raw output. external_axes: Optional[List[plt.axes]] @@ -167,47 +160,48 @@ def show_treasuries( ---------- Plots the Treasury Series. """ - if external_axes is None: - _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI) - else: - ax = external_axes[0] treasury_data = econdb_model.get_treasuries( types, maturities, frequency, start_date, end_date ) - for treasury, maturities_data in treasury_data.items(): - for maturity in maturities_data: - ax.plot(maturities_data[maturity], label=f"{treasury} [{maturity}]") + if not store: + if external_axes is None: + _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI) + else: + ax = external_axes[0] - ax.set_title("U.S. Treasuries") - ax.legend() + for treasury, maturities_data in treasury_data.items(): + for maturity in maturities_data: + ax.plot(maturities_data[maturity], label=f"{treasury} [{maturity}]") - if raw or export: - df = pd.DataFrame.from_dict(treasury_data, orient="index").stack().to_frame() - df = pd.DataFrame(df[0].values.tolist(), index=df.index).T - df.columns = ["_".join(column) for column in df.columns] + ax.set_title("U.S. Treasuries") + ax.legend() - if raw: - print_rich_table( - df.iloc[-10:], - headers=list(df.columns), - show_index=True, - title="U.S. Treasuries", - ) + theme.style_primary_axis(ax) - if export: - export_data( - export, - os.path.dirname(os.path.abspath(__file__)), - "treasuries_data", - df, - ) + if external_axes is None: + theme.visualize_output() - theme.style_primary_axis(ax) + df = pd.DataFrame.from_dict(treasury_data, orient="index").stack().to_frame() + df = pd.DataFrame(df[0].values.tolist(), index=df.index).T + df.columns = ["_".join(column) for column in df.columns] - if external_axes is None: - theme.visualize_output() + if raw: + print_rich_table( + df.iloc[-10:], + headers=list(df.columns), + show_index=True, + title="U.S. Treasuries", + ) + + if export: + export_data( + export, + os.path.dirname(os.path.abspath(__file__)), + "treasuries_data", + df, + ) @log_start_end(log=logger) diff --git a/gamestonk_terminal/economy/economy_controller.py b/gamestonk_terminal/economy/economy_controller.py index 3e3da48f78d..77ea90b8e46 100644 --- a/gamestonk_terminal/economy/economy_controller.py +++ b/gamestonk_terminal/economy/economy_controller.py @@ -6,7 +6,7 @@ import argparse import logging import os from datetime import date -from typing import List, Dict +from typing import List, Dict, Any import pandas as pd from prompt_toolkit.completion import NestedCompleter @@ -26,6 +26,7 @@ from gamestonk_terminal.economy import ( fred_model, yfinance_model, yfinance_view, + plot_view, ) from gamestonk_terminal.helper_funcs import ( EXPORT_BOTH_RAW_DATA_AND_FIGURES, @@ -52,6 +53,7 @@ class EconomyController(BaseController): "fred", "index", "treasury", + "plot", "valuation", "performance", "spectrum", @@ -92,6 +94,7 @@ class EconomyController(BaseController): "unemp", ] overview_options = ["indices", "usbonds", "glbonds", "currencies"] + futures_options = ["energy", "metals", "meats", "grains", "softs"] tyld_maturity = ["3m", "5y", "10y", "30y"] valuation_sort_cols = [ "Name", @@ -148,43 +151,69 @@ class EconomyController(BaseController): super().__init__(queue) self.current_series: Dict = dict() + self.fred_query: pd.Series = pd.Series() + self.DATASETS: Dict[Any, pd.DataFrame] = dict() if session and gtff.USE_PROMPT_TOOLKIT: - choices: dict = {c: {} for c in self.controller_choices} - choices["overview"] = {c: None for c in self.overview_options} + self.choices: dict = {c: {} for c in self.controller_choices} + self.choices["overview"] = {c: None for c in self.overview_options} - choices["futures"] = { - c: None for c in ["energy", "metals", "meats", "grains", "softs"] + self.choices["futures"] = {c: None for c in self.futures_options} + + self.choices["index"] = {c: None for c in yfinance_model.INDICES} + + self.choices["macro"]["-p"] = {c: None for c in econdb_model.PARAMETERS} + self.choices["macro"]["--parameter"] = { + c: None for c in econdb_model.PARAMETERS } - - choices["index"] = {c: None for c in yfinance_model.INDICES} - - choices["macro"]["-p"] = {c: None for c in econdb_model.PARAMETERS} - choices["macro"]["--parameter"] = {c: None for c in econdb_model.PARAMETERS} - choices["macro"]["-c"] = {c: None for c in econdb_model.COUNTRY_CODES} - choices["macro"]["--countries"] = { + self.choices["macro"]["-c"] = {c: None for c in econdb_model.COUNTRY_CODES} + self.choices["macro"]["--countries"] = { c: None for c in econdb_model.COUNTRY_CODES } - choices["valuation"]["-s"] = {c: None for c in self.valuation_sort_cols} - choices["valuation"]["--sortby"] = { + self.choices["valuation"]["-s"] = { + c: None for c in self.valuation_sort_cols + } + self.choices["valuation"]["--sortby"] = { c: None for c in self.valuation_sort_cols } - choices["performance"]["-s"] = {c: None for c in self.performance_sort_list} - choices["performance"]["--sortby"] = { + self.choices["performance"]["-s"] = { + c: None for c in self.performance_sort_list + } + self.choices["performance"]["--sortby"] = { c: None for c in self.performance_sort_list } - choices["map"]["-p"] = {c: None for c in self.map_period_list} - choices["map"]["--period"] = {c: None for c in self.map_period_list} + self.choices["map"]["-p"] = {c: None for c in self.map_period_list} + self.choices["map"]["--period"] = {c: None for c in self.map_period_list} - choices["feargreed"]["-i"] = {c: None for c in self.fear_greed_indicators} - choices["feargreed"]["--indicator"] = { + self.choices["feargreed"]["-i"] = { + c: None for c in self.fear_greed_indicators + } + self.choices["feargreed"]["--indicator"] = { c: None for c in self.fear_greed_indicators } - self.completer = NestedCompleter.from_nested_dict(choices) + self.completer = NestedCompleter.from_nested_dict(self.choices) + + def update_runtime_choices(self): + if session and gtff.USE_PROMPT_TOOLKIT: + if not self.fred_query.empty: + self.choices["fred"] = {c: None for c in self.fred_query} + if self.DATASETS: + options = [ + option + for _, values in self.DATASETS.items() + for option in values.keys() + ] + + for argument in ["-y1", "-yaxis1", "-y2", "yaxis2"]: + self.choices["plot"][argument] = { + option: None for option in options + } + + self.completer = NestedCompleter.from_nested_dict(self.choices) def print_help(self): """Print help""" @@ -199,6 +228,7 @@ Macro Data fred collect macro data from FRED based on a series ID [src][Source: FRED][/src] index find and plot any (major) index on the market [src][Source: Yahoo Finance][/src] treasury obtain U.S. treasury rates [src][Source: EconDB][/src] + plot plot data from the above commands together Performance & Valuations rtps real-time performance sectors [src][Source: Alpha Vantage][/src] @@ -361,7 +391,10 @@ Index add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter, prog="macro", - description="Get a broad selection of macro data from one or multiple countries. [Source: EconDB]", + description="Get a broad selection of macro data from one or multiple countries. This includes Gross " + "Domestic Product (RGDP & GDP) and the underlying components, Treasury Yields (Y10YD & M3YD), " + "Employment figures (URATE, EMP, AC0I0 and EMRATIO), Government components (e.g. GBAL & GREV), " + "Consumer and Producer Indices (CPI & PPI) and a variety of other indicators. [Source: EconDB]", ) parser.add_argument( @@ -421,10 +454,19 @@ Index "--convert_currency", dest="convert_currency", help="Convert the currency of the chosen country to a specified currency. By default, this will be USD " - "unless specified with this command. To find the currencies use the option -sc", + "unless specified with this command. To find the currency symbols use the option -sc", default="USD", ) + parser.add_argument( + "-st", + "--store", + dest="store", + help="Store the data to be used for plotting with the 'plot' command.", + action="store_true", + default=False, + ) + ns_parser = parse_known_args_and_warn( parser, other_args, export_allowed=EXPORT_ONLY_RAW_DATA_ALLOWED, raw=True ) @@ -442,15 +484,30 @@ Index headers=["Country", "Currency"], ) elif ns_parser.parameters and ns_parser.countries: - econdb_view.show_macro_data( - parameters=ns_parser.parameters, - countries=ns_parser.countries, - start_date=ns_parser.start_date, - end_date=ns_parser.end_date, - convert_currency=ns_parser.convert_currency, - raw=ns_parser.raw, - export=ns_parser.export, - ) + if ns_parser.store: + self.DATASETS["macro"] = econdb_model.get_aggregated_macro_data( + parameters=ns_parser.parameters, + countries=ns_parser.countries, + start_date=ns_parser.start_date, + end_date=ns_parser.end_date, + convert_currency=ns_parser.convert_currency, + ) + + self.DATASETS["macro"].columns = [ + "_".join(column) for column in self.DATASETS["macro"].columns + ] + else: + econdb_view.show_macro_data( + parameters=ns_parser.parameters, + countries=ns_parser.countries, + start_date=ns_parser.start_date, + end_date=ns_parser.end_date, + convert_currency=ns_parser.convert_currency, + raw=ns_parser.raw, + export=ns_parser.export, + ) + + self.update_runtime_choices() def call_fred(self, other_args: List[str]): """Process fred command""" @@ -471,21 +528,40 @@ Index parser.add_argument( "-s", + "--start_date", dest="start_date", type=valid_date, - default="2020-01-01", help="Starting date (YYYY-MM-DD) of data", + default=None, + ) + + parser.add_argument( + "-e", + "--end_date", + dest="end_date", + type=valid_date, + help="Ending date (YYYY-MM-DD) of data", + default=None, ) parser.add_argument( "-q", "--query", + nargs="+", action="store", dest="query", - type=str, help="Query the FRED database to obtain Series IDs given the query seaarch term.", ) + parser.add_argument( + "-st", + "--store", + dest="store", + help="Store the data to be used for plotting with the 'plot' command.", + action="store_true", + default=False, + ) + if other_args and "-" not in other_args[0][0]: other_args.insert(0, "-p") ns_parser = parse_known_args_and_warn( @@ -497,10 +573,13 @@ Index ) if ns_parser: if ns_parser.query: - fred_view.notes( - series_term=ns_parser.query, - num=ns_parser.limit, - ) + query = " ".join(ns_parser.query) + df_search = fred_model.get_series_notes(query) + fred_view.notes(series_term=query, num=ns_parser.limit) + + self.fred_query = df_search["id"].head(ns_parser.limit) + self.update_runtime_choices() + if ns_parser.parameter: series_dict = {} for series in ns_parser.parameter: @@ -512,15 +591,23 @@ Index "units": information["seriess"][0]["units_short"], } - self.current_series = {series: series_dict[series]} + self.current_series = {series: series_dict[series]} - fred_view.display_fred_series( - series_dict, - ns_parser.start_date, - ns_parser.raw, - ns_parser.export, - ns_parser.limit, - ) + if ns_parser.store: + self.DATASETS["fred"] = fred_model.get_aggregated_series_data( + series_dict, ns_parser.start_date, ns_parser.end_date + ) + else: + fred_view.display_fred_series( + d_series=series_dict, + start_date=ns_parser.start_date, + end_date=ns_parser.end_date, + raw=ns_parser.raw, + export=ns_parser.export, + limit=ns_parser.limit, + ) + + self.update_runtime_choices() @log_start_end(log=logger) def call_index(self, other_args: List[str]): @@ -610,6 +697,15 @@ Index help="Search for indices with given keyword", ) + parser.add_argument( + "-st", + "--store", + dest="store", + help="Store the data to be used for plotting with the 'plot' command.", + action="store_true", + default=False, + ) + if other_args and "-" not in other_args[0][0]: other_args.insert(0, "-i") ns_parser = parse_known_args_and_warn( @@ -631,15 +727,29 @@ Index title="Major Indices", ) elif ns_parser.indices: - yfinance_view.show_indices( - indices=ns_parser.indices, - interval=ns_parser.interval, - start_date=ns_parser.start_date, - end_date=ns_parser.end_date, - column=ns_parser.column, - raw=ns_parser.raw, - export=ns_parser.export, - ) + if ns_parser.store: + self.DATASETS["index"] = pd.DataFrame() + for index in ns_parser.indices: + self.DATASETS["index"][index] = yfinance_model.get_index( + index, + interval=ns_parser.interval, + start_date=ns_parser.start_date, + end_date=ns_parser.end_date, + column=ns_parser.column, + ) + else: + yfinance_view.show_indices( + indices=ns_parser.indices, + interval=ns_parser.interval, + start_date=ns_parser.start_date, + end_date=ns_parser.end_date, + column=ns_parser.column, + store=ns_parser.store, + raw=ns_parser.raw, + export=ns_parser.export, + ) + + self.update_runtime_choices() @log_start_end(log=logger) def call_treasury(self, other_args: List[str]): @@ -650,7 +760,8 @@ Index prog="treasury", description="Obtain any set of U.S. treasuries and plot them together. These can be a range of maturities " "for nominal, inflation-adjusted (on long term average of inflation adjusted) and secondary " - "markets over a lengthy period. [Source: EconDB / FED]", + "markets over a lengthy period. Note: 3-month and 10-year treasury yields for other countries " + "are available via the command 'macro' and parameter 'Y10YD' and 'M3YD'. [Source: EconDB / FED]", ) parser.add_argument( @@ -710,6 +821,15 @@ Index default=date.today(), ) + parser.add_argument( + "-st", + "--store", + dest="store", + help="Store the data to be used for plotting with the 'plot' command.", + action="store_true", + default=False, + ) + if other_args and "-" not in other_args[0][0]: other_args.insert(0, "-m") ns_parser = parse_known_args_and_warn( @@ -724,15 +844,141 @@ Index if ns_parser.show_maturities: econdb_view.show_treasury_maturities(econdb_model.TREASURIES) elif ns_parser.maturity and ns_parser.type: - econdb_view.show_treasuries( - types=ns_parser.type, - maturities=ns_parser.maturity, - frequency=ns_parser.frequency, - start_date=ns_parser.start_date, - end_date=ns_parser.end_date, - raw=ns_parser.raw, - export=ns_parser.export, + if ns_parser.store: + treasury_data = econdb_model.get_treasuries( + instruments=ns_parser.type, + maturities=ns_parser.maturity, + frequency=ns_parser.frequency, + start_date=ns_parser.start_date, + end_date=ns_parser.end_date, + ) + + df = ( + pd.DataFrame.from_dict(treasury_data, orient="index") + .stack() + .to_frame() + ) + self.DATASETS["treasuries"] = pd.DataFrame( + df[0].values.tolist(), index=df.index + ).T + self.DATASETS["treasuries"].columns = [ + "_".join(column) + for column in self.DATASETS["treasuries"].columns + ] + else: + econdb_view.show_treasuries( + types=ns_parser.type, + maturities=ns_parser.maturity, + frequency=ns_parser.frequency, + start_date=ns_parser.start_date, + end_date=ns_parser.end_date, + store=ns_parser.store, + raw=ns_parser.raw, + export=ns_parser.export, + ) + + self.update_runtime_choices() + + @log_start_end(log=logger) + def call_plot(self, other_args: List[str]): + """Process plot command""" + parser = argparse.ArgumentParser( + add_help=False, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + prog="plot", + description="Obtain any set of U.S. treasuries and plot them together. These can be a range of maturities " + "for nominal, inflation-adjusted (on long term average of inflation adjusted) and secondary " + "markets over a lengthy period. Note: 3-month and 10-year treasury yields for other countries " + "are available via the command 'macro' and parameter 'Y10YD' and 'M3YD'. [Source: EconDB / FED]", + ) + + parser.add_argument( + "-y1", + "--yaxis_1", + nargs="+", + dest="yaxis_1", + help="Select the data you wish to plot on the first y-axis. You can select multiple variables here.", + default="", + ) + + parser.add_argument( + "-y2", + "--yaxis_2", + nargs="+", + dest="yaxis_2", + help="Select the data you wish to plot on the second y-axis. You can select multiple variables here.", + default="", + ) + + parser.add_argument( + "-s", + "--show", + dest="show", + help="Show all available options currently collected.", + action="store_true", + default=False, + ) + + if other_args and "-" not in other_args[0][0]: + other_args.insert(0, "-y1") + ns_parser = parse_known_args_and_warn( + parser, + other_args, + export_allowed=EXPORT_ONLY_RAW_DATA_ALLOWED, + raw=True, + limit=10, + ) + + if ns_parser: + if ns_parser.show: + options = { + command: ", ".join(values.keys()) + for command, values in self.DATASETS.items() + } + print_rich_table( + pd.DataFrame.from_dict( + options, orient="index", columns=["Options"] + ), + show_index=True, + index_name="Command", + title="Options available to plot", ) + else: + dataset_yaxis_1 = pd.DataFrame() + dataset_yaxis_2 = pd.DataFrame() + + if ns_parser.yaxis_1: + for variable in ns_parser.yaxis_1: + for _, data in self.DATASETS.items(): + if variable in data.columns: + dataset_yaxis_1[variable] = data[variable] + break + if variable not in dataset_yaxis_1: + return console.print( + f"Not able to find any data for the -y1 argument {ns_parser.yaxis_1}. " + f"The currently available options are: " + f"{', '.join(self.choices['-y1'])}" + ) + + if ns_parser.yaxis_2: + for variable in ns_parser.yaxis_2: + for _, data in self.DATASETS.items(): + if variable in data.columns: + dataset_yaxis_2[variable] = data[variable] + break + if variable not in dataset_yaxis_2: + return console.print( + f"Not able to find any data for the -y2 argument {ns_parser.yaxis_2}. " + f"The currently available options are: " + f"{', '.join(self.choices['-y2'])}" + ) + + if ns_parser.yaxis_1 or ns_parser.yaxis_2: + return plot_view.show_plot( + dataset_yaxis_1=dataset_yaxis_1, dataset_yaxis_2=dataset_yaxis_2 + ) + + console.print() @log_start_end(log=logger) def call_map(self, other_args: List[str]): diff --git a/gamestonk_terminal/economy/fred_model.py b/gamestonk_terminal/economy/fred_model.py index 21157a03266..9a443706463 100644 --- a/gamestonk_terminal/economy/fred_model.py +++ b/gamestonk_terminal/economy/fred_model.py @@ -136,7 +136,7 @@ def get_series_ids(series_term: str, num: int) -> Tuple[List[str], List[str]]: @log_start_end(log=logger) -def get_series_data(series_id: str, start: str) -> pd.DataFrame: +def get_series_data(series_id: str, start: str = None, end: str = None) -> pd.DataFrame: """Get Series data. [Source: FRED] Parameters ---------- @@ -144,6 +144,9 @@ def get_series_data(series_id: str, start: str) -> pd.DataFrame: Series ID to get data from start : str Start date to get data from, format yyyy-mm-dd + end : str + End data to get from, format yyyy-mm-dd + Returns ---------- pd.DataFrame @@ -153,9 +156,43 @@ def get_series_data(series_id: str, start: str) -> pd.DataFrame: try: fredapi_client = Fred(cfg.API_FRED_KEY) - df = fredapi_client.get_series(series_id, start) + df = fredapi_client.get_series(series_id, start, end) # Series does not exist & invalid api keys except HTTPError as e: console.print(e) return df + + +@log_start_end(log=logger) +def get_aggregated_series_data( + d_series: Dict[str, Dict[str, str]], start: str = None, end: str = None +) -> pd.DataFrame: + """Get Series data. [Source: FRED] + Parameters + ---------- + series_id : str + Series ID to get data from + start : str + Start date to get data from, format yyyy-mm-dd + end : str + End data to get from, format yyyy-mm-dd + + Returns + ---------- + pd.DataFrame + Series data + """ + series_ids = list(d_series.keys()) + data = pd.DataFrame() + + for s_id in series_ids: + data = pd.concat( + [ + data, + pd.DataFrame(get_series_data(s_id, start, end), columns=[s_id]), + ], + axis=1, + ) + + return data diff --git a/gamestonk_terminal/economy/fred_view.py b/gamestonk_terminal/economy/fred_view.py index 6ada2e2fecb..a7020267fa6 100644 --- a/gamestonk_terminal/economy/fred_view.py +++ b/gamestonk_terminal/economy/fred_view.py @@ -11,9 +11,9 @@ import numpy as np import pandas as pd from pandas.plotting import register_matplotlib_converters -from gamestonk_terminal.decorators import check_api_key -from gamestonk_terminal.config_terminal import theme from gamestonk_terminal.config_plot import PLOT_DPI +from gamestonk_terminal.config_terminal import theme +from gamestonk_terminal.decorators import check_api_key from gamestonk_terminal.decorators import log_start_end from gamestonk_terminal.economy import fred_model from gamestonk_terminal.helper_funcs import ( @@ -47,8 +47,9 @@ def format_units(num: int) -> str: @log_start_end(log=logger) @check_api_key(["API_FRED_KEY"]) -def notes(series_term: str, num: int): +def notes(series_term: str, num: int) -> pd.DataFrame: """Print Series notes. [Source: FRED] + Parameters ---------- series_term : str @@ -80,6 +81,7 @@ def notes(series_term: str, num: int): def display_fred_series( d_series: Dict[str, Dict[str, str]], start_date: str, + end_date: str = "", raw: bool = False, export: str = "", limit: int = 10, @@ -89,10 +91,14 @@ def display_fred_series( Parameters ---------- - series : str + d_series : str FRED Series ID from https://fred.stlouisfed.org. For multiple series use: series1,series2,series3 start_date : str Starting date (YYYY-MM-DD) of data + end_date : str + Ending date (YYYY-MM-DD) of data + store : bool + Whether to prevent plotting the data. raw : bool Output only raw data export : str @@ -103,61 +109,53 @@ def display_fred_series( External axes (3 axes are expected in the list), by default None """ series_ids = list(d_series.keys()) - data = pd.DataFrame() + data = fred_model.get_aggregated_series_data(d_series, start_date, end_date) - for s_id in series_ids: - data = pd.concat( - [ - data, - pd.DataFrame( - fred_model.get_series_data(s_id, start_date), columns=[s_id] - ), - ], - axis=1, - ) - # Try to get everything onto the same 0-10 scale. - # To do so, think in scientific notation. Divide the data by whatever the E would be - if external_axes is None: - _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI) - - else: - if len(external_axes) != 3: - logger.error("Expected list of 3 axis items") - console.print("[red]Expected list of 3 axis items./n[/red]") - return - (ax,) = external_axes - - if len(series_ids) == 1: - s_id = series_ids[0] - sub_dict: Dict = d_series[s_id] - title = f"{sub_dict['title']} ({sub_dict['units']})" - ax.plot(data.index, data, label="\n".join(textwrap.wrap(title, 80))) - else: - for s_id, sub_dict in d_series.items(): - data_to_plot = data[s_id].dropna() - exponent = int(np.log10(data_to_plot.max())) - data_to_plot /= 10**exponent - multiplier = f"x {format_units(10**exponent)}" if exponent > 0 else "" - title = f"{sub_dict['title']} ({sub_dict['units']}) {'['+multiplier+']' if multiplier else ''}" - ax.plot( - data_to_plot.index, - data_to_plot, - label="\n".join(textwrap.wrap(title, 80)) - if len(series_ids) < 5 - else title, - ) - - ax.legend(prop={"size": 10}, bbox_to_anchor=(0, 1), loc="lower left") if data.empty: logger.error("No data") - console.print("[red]No data[/red]\n") + console.print("[red]No data available.[/red]\n") else: + # Try to get everything onto the same 0-10 scale. + # To do so, think in scientific notation. Divide the data by whatever the E would be + if external_axes is None: + _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI) + + else: + if len(external_axes) != 3: + logger.error("Expected list of 3 axis items") + console.print("[red]Expected list of 3 axis items./n[/red]") + return + (ax,) = external_axes + + if len(series_ids) == 1: + s_id = series_ids[0] + sub_dict: Dict = d_series[s_id] + title = f"{sub_dict['title']} ({sub_dict['units']})" + ax.plot(data.index, data, label="\n".join(textwrap.wrap(title, 80))) + else: + for s_id, sub_dict in d_series.items(): + data_to_plot = data[s_id].dropna() + exponent = int(np.log10(data_to_plot.max())) + data_to_plot /= 10**exponent + multiplier = f"x {format_units(10**exponent)}" if exponent > 0 else "" + title = f"{sub_dict['title']} ({sub_dict['units']}) {'['+multiplier+']' if multiplier else ''}" + ax.plot( + data_to_plot.index, + data_to_plot, + label="\n".join(textwrap.wrap(title, 80)) + if len(series_ids) < 5 + else title, + ) + + ax.legend(prop={"size": 10}, bbox_to_anchor=(0, 1), loc="lower left") + ax.set_xlim(data.index[0], data.index[-1]) theme.style_primary_axis(ax) if external_axes is None: theme.visualize_output() data.index = [x.strftime("%Y-%m-%d") for x in data.index] + if raw: print_rich_table( data.tail(limit), @@ -166,9 +164,10 @@ def display_fred_series( index_name="Date", ) console.print("") + export_data( export, os.path.dirname(os.path.abspath(__file__)), - "plot", + "fred", data, ) diff --git a/gamestonk_terminal/economy/plot_view.py b/gamestonk_terminal/economy/plot_view.py new file mode 100644 index 00000000000..82d09aaa7cb --- /dev/null +++ b/gamestonk_terminal/economy/plot_view.py @@ -0,0 +1,69 @@ +""" Plot Controller """ +from typing import Optional, List + +from matplotlib import pyplot as plt + +from gamestonk_terminal.config_plot import PLOT_DPI +from gamestonk_terminal.config_terminal import theme +from gamestonk_terminal.helper_funcs import plot_autoscale + + +def show_plot( + dataset_yaxis_1, dataset_yaxis_2, external_axes: Optional[List[plt.Axes]] = None +): + """ + The ability to plot any data coming from EconDB, FRED or Yahoo Finance. + + Parameters + ---------- + dataset_yaxis_1: pd.DataFrame + Data you wish to plot on the first y-axis. + dataset_yaxis_2 : pd.DataFrame + Data you wish to plot on the second y-axis. + external_axes: Optional[List[plt.axes]] + External axes to plot on + + Returns + ---------- + Plots the data. + """ + if external_axes is None: + _, ax1 = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI) + + else: + ax1, ax2 = external_axes + + color_palette = theme.get_colors() + ax_1_coloring = 0 + ax_2_coloring = -1 + + for column in dataset_yaxis_1: + ax1.plot( + dataset_yaxis_1[column], + label=column.replace("_", " "), + color=color_palette[ax_1_coloring], + ) + ax_1_coloring += 1 + + theme.style_primary_axis(ax1) + ax1.yaxis.set_label_position("left") + ax1.legend() + + if not dataset_yaxis_2.empty: + ax2 = ax1.twinx() + + for column in dataset_yaxis_2: + ax2.plot( + dataset_yaxis_2[column], + label=column, + color=color_palette[ax_2_coloring], + ) + ax_2_coloring += -1 + + ax2.yaxis.set_label_position("right") + theme.style_twin_axis(ax2) + + ax2.legend(loc="upper right") + + if external_axes is None: + theme.visualize_output() diff --git a/gamestonk_terminal/economy/yfinance_view.py b/gamestonk_terminal/economy/yfinance_view.py index f6facab740d..4b85b2cb51b 100644 --- a/gamestonk_terminal/economy/yfinance_view.py +++ b/gamestonk_terminal/economy/yfinance_view.py @@ -1,5 +1,6 @@ """ EconDB View """ __docformat__ = "numpy" +# pylint:disable=too-many-arguments import logging import os @@ -29,6 +30,7 @@ def show_indices( start_date: int = None, end_date: int = None, column: str = "Adj Close", + store: bool = False, raw: bool = False, external_axes: Optional[List[plt.axes]] = None, export: str = "", @@ -46,6 +48,10 @@ def show_indices( The starting date, format "YEAR-MONTH-DAY", i.e. 2010-12-31. end_date : str The end date, format "YEAR-MONTH-DAY", i.e. 2020-06-05. + column : str + Which column to load in, by default this is the Adjusted Close. + store : bool + Whether to prevent plotting the data. raw : bool Whether to display the raw output. external_axes: Optional[List[plt.axes]] @@ -59,24 +65,31 @@ def show_indices( """ indices_data: pd.DataFrame = pd.DataFrame() - if external_axes is None: - _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI) - else: - ax = external_axes[0] - for index in indices: indices_data[index] = get_index(index, interval, start_date, end_date, column) - if index in INDICES: - label = INDICES[index]["name"] + if not store: + if external_axes is None: + _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI) else: - label = index + ax = external_axes[0] - if not indices_data[index].empty: - ax.plot(indices_data[index], label=label) + for index in indices: + if index in INDICES: + label = INDICES[index]["name"] + else: + label = index - ax.set_title("Indices") - ax.legend() + if not indices_data[index].empty: + ax.plot(indices_data[index], label=label) + + ax.set_title("Indices") + ax.legend() + + theme.style_primary_axis(ax) + + if external_axes is None: + theme.visualize_output() if raw: print_rich_table( @@ -94,10 +107,7 @@ def show_indices( indices_data, ) - theme.style_primary_axis(ax) - - if external_axes is None: - theme.visualize_output() + return indices_data @log_start_end(log=logger)