Fix economy plot (#3346)

* function to improve the data to plot

* added more user messages

* changed view to be able to return data - so the request is not repeated

* moved ax1.lengend up

* removes nan values from dataframe

* fixed tests on the economy controller

* fixes economy notebook

* removed unsuded Dict

* function does not return anything by default

* function does not return anything by default

* Revert "function does not return anything by default"

This reverts commit 8aca895a5e.

* function does not return anything by default

* removed wrong return type

* lint

openbb_terminal/economy/economy_controller.py

@@ -810,6 +810,7 @@ class EconomyController(BaseController):
         )
         if ns_parser:
             parameters = list_from_str(ns_parser.parameter.upper())
+
             if ns_parser.query:
                 query = ns_parser.query.replace(",", " ")
                 df_search = fred_model.get_series_notes(search_query=query)
@@ -820,6 +821,11 @@ class EconomyController(BaseController):
                     self.fred_query = df_search["id"].head(ns_parser.limit)
                     self.update_runtime_choices()
 
+                if parameters:
+                    console.print(
+                        "\nWarning: -p/--parameter is ignored when using -q/--query."
+                    )
+
                 return self.queue
 
             if parameters:
@@ -838,40 +844,42 @@ class EconomyController(BaseController):
                 if not series_dict:
                     return self.queue
 
-                df, detail = fred_model.get_aggregated_series_data(
+                df, detail = fred_view.display_fred_series(
                     series_ids=parameters,
                     start_date=ns_parser.start_date,
                     end_date=ns_parser.end_date,
+                    limit=ns_parser.limit,
+                    raw=ns_parser.raw,
+                    export=ns_parser.export,
+                    get_data=True,
                 )
 
-                for series_id, data in detail.items():
-                    self.FRED_TITLES[series_id] = f"{data['title']} ({data['units']})"
-
                 if not df.empty:
-                    self.DATASETS["fred"] = pd.concat(
-                        [
-                            self.DATASETS["fred"],
-                            df,
-                        ]
-                    )
+
+                    for series_id, data in detail.items():
+                        self.FRED_TITLES[
+                            series_id
+                        ] = f"{data['title']} ({data['units']})"
+
+                        # Making data available at the class level
+                        self.DATASETS["fred"][series_id] = df[series_id]
 
                     self.stored_datasets = (
                         economy_helpers.update_stored_datasets_string(self.DATASETS)
                     )
 
-                    fred_view.display_fred_series(
-                        series_ids=parameters,
-                        start_date=ns_parser.start_date,
-                        end_date=ns_parser.end_date,
-                        limit=ns_parser.limit,
-                        raw=ns_parser.raw,
-                        export=ns_parser.export,
-                    )
-
                     self.update_runtime_choices()
                     if obbff.ENABLE_EXIT_AUTO_HELP:
                         self.print_help()
 
+                else:
+                    console.print("[red]No data found for the given Series ID[/red]")
+
+            elif not parameters and ns_parser.raw:
+                console.print(
+                    "Warning: -r/--raw should be combined with -p/--parameter."
+                )
+
     @log_start_end(log=logger)
     def call_index(self, other_args: List[str]):
         """Process index command"""
@@ -1441,9 +1449,17 @@ class EconomyController(BaseController):
                                         f"{country}{transformtype}[{parameter}, Units: {units}]"
                                     ] = data[variable]
                                 elif key == "fred":
-                                    dataset_yaxis1[self.FRED_TITLES[variable]] = data[
-                                        variable
-                                    ]
+                                    compound_detail = self.FRED_TITLES[variable]
+                                    detail = {
+                                        "units": compound_detail.split("(")[-1].split(
+                                            ")"
+                                        )[0],
+                                        "title": compound_detail.split("(")[0].strip(),
+                                    }
+                                    data_to_plot, title = fred_view.format_data_to_plot(
+                                        data[variable], detail
+                                    )
+                                    dataset_yaxis1[title] = data_to_plot
                                 elif (
                                     key == "index"
                                     and variable in yfinance_model.INDICES
@@ -1503,9 +1519,17 @@ class EconomyController(BaseController):
                                         f"{country}{transformtype}[{parameter}, Units: {units}]"
                                     ] = data[variable]
                                 elif key == "fred":
-                                    dataset_yaxis2[self.FRED_TITLES[variable]] = data[
-                                        variable
-                                    ]
+                                    compound_detail = self.FRED_TITLES[variable]
+                                    detail = {
+                                        "units": compound_detail.split("(")[-1].split(
+                                            ")"
+                                        )[0],
+                                        "title": compound_detail.split("(")[0].strip(),
+                                    }
+                                    data_to_plot, title = fred_view.format_data_to_plot(
+                                        data[variable], detail
+                                    )
+                                    dataset_yaxis2[title] = data_to_plot
                                 elif (
                                     key == "index"
                                     and variable in yfinance_model.INDICES

openbb_terminal/economy/fred_model.py

@@ -204,7 +204,7 @@ def get_series_data(
 @check_api_key(["API_FRED_KEY"])
 def get_aggregated_series_data(
     series_ids: List[str], start_date: str = None, end_date: str = None
-) -> pd.DataFrame:
+) -> Tuple[pd.DataFrame, dict]:
     """Get Series data. [Source: FRED]
     Parameters
     ----------
@@ -219,6 +219,8 @@ def get_aggregated_series_data(
     ----------
     pd.DataFrame
         Series data
+    dict
+        Dictionary of series ids and titles
     """
 
     data = pd.DataFrame()
@@ -234,15 +236,12 @@ def get_aggregated_series_data(
             }
 
     for s_id in series_ids:
-        data = pd.concat(
-            [
-                data,
-                pd.DataFrame(
-                    get_series_data(s_id, start_date, end_date), columns=[s_id]
-                ),
-            ],
-            axis=1,
-        )
+
+        series = pd.DataFrame(
+            get_series_data(s_id, start_date, end_date), columns=[s_id]
+        ).dropna()
+
+        data[s_id] = series[s_id]
 
     return data, detail
 

openbb_terminal/economy/fred_view.py

@@ -4,7 +4,7 @@ __docformat__ = "numpy"
 import logging
 import os
 import textwrap
-from typing import Dict, Optional, List
+from typing import Optional, List, Tuple
 from datetime import datetime
 
 import matplotlib.pyplot as plt
@@ -49,7 +49,7 @@ def format_units(num: int) -> str:
 
 @log_start_end(log=logger)
 @check_api_key(["API_FRED_KEY"])
-def notes(search_query: str, limit: int = 10) -> pd.DataFrame:
+def notes(search_query: str, limit: int = 10):
     """Display series notes. [Source: FRED]
 
     Parameters
@@ -79,6 +79,7 @@ def display_fred_series(
     start_date: str = None,
     end_date: str = None,
     limit: int = 10,
+    get_data: bool = False,
     raw: bool = False,
     export: str = "",
     external_axes: Optional[List[plt.Axes]] = None,
@@ -118,29 +119,19 @@ def display_fred_series(
         elif is_valid_axes_count(external_axes, 1):
             (ax,) = external_axes
         else:
-            return
+            return None
+
+        for s_id, sub_dict in detail.items():
+
+            data_to_plot, title = format_data_to_plot(data[s_id], sub_dict)
 
-        if len(series_ids) == 1:
-            s_id = series_ids[0]
-            sub_dict: Dict = detail[s_id]
-            title = f"{sub_dict['title']} ({sub_dict['units']})"
             ax.plot(
-                data.index, data.iloc[:, 0], label="\n".join(textwrap.wrap(title, 80))
+                data_to_plot.index,
+                data_to_plot,
+                label="\n".join(textwrap.wrap(title, 80))
+                if len(series_ids) < 5
+                else title,
             )
-        else:
-            for s_id, sub_dict in detail.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(
             bbox_to_anchor=(0, 0.40, 1, -0.52),
@@ -172,6 +163,25 @@ def display_fred_series(
             data,
         )
 
+    if get_data:
+        return data, detail
+
+    return None
+
+
+def format_data_to_plot(data: pd.DataFrame, detail: dict) -> Tuple[pd.DataFrame, str]:
+    """Helper to format data to plot"""
+
+    data_to_plot = data.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"{detail['title']} ({detail['units']}) {'['+multiplier+']' if multiplier else ''}"
+
+    data_to_plot.index = pd.to_datetime(data_to_plot.index)
+
+    return data_to_plot, title
+
 
 @log_start_end(log=logger)
 @check_api_key(["API_FRED_KEY"])

openbb_terminal/economy/plot_view.py

@@ -47,6 +47,8 @@ def show_plot(
     ax_1_coloring = 0
     ax_2_coloring = -1
 
+    dataset_yaxis_1 = dataset_yaxis_1.dropna()
+
     for column in dataset_yaxis_1:
         ax1.plot(
             dataset_yaxis_1[column],
@@ -55,11 +57,22 @@ def show_plot(
         )
         ax_1_coloring += 1
 
+    ax1.legend(
+        [fill(column, 45) for column in dataset_yaxis_1.columns],
+        bbox_to_anchor=(0, 0.40, 1, -0.52),
+        loc="upper right",
+        mode="expand",
+        borderaxespad=0,
+        prop={"size": 9},
+    )
+
     theme.style_primary_axis(ax1)
 
     if not dataset_yaxis_2.empty:
         ax2 = ax1.twinx()
 
+        dataset_yaxis_2 = dataset_yaxis_2.dropna()
+
         for column in dataset_yaxis_2:
             ax2.plot(
                 dataset_yaxis_2[column],
@@ -79,15 +92,6 @@ def show_plot(
             prop={"size": 9},
         )
 
-    ax1.legend(
-        [fill(column, 45) for column in dataset_yaxis_1.columns],
-        bbox_to_anchor=(0, 0.40, 1, -0.52),
-        loc="upper right",
-        mode="expand",
-        borderaxespad=0,
-        prop={"size": 9},
-    )
-
     if external_axes is None:
         theme.visualize_output()
 

openbb_terminal/reports/templates/economy.ipynb

@@ -32,7 +32,7 @@
     "\n",
     "from openbb_terminal.reports import widget_helpers as widgets\n",
     "from openbb_terminal.sdk import openbb\n",
-    "from openbb_terminal.sdk import openbb import helper\n",
+    "from openbb_terminal.helper_funcs import get_user_timezone\n",
     "from openbb_terminal.helper_classes import TerminalStyle\n",
     "from openbb_terminal.core.config.paths import REPOSITORY_DIRECTORY\n",
     "\n",
@@ -83,7 +83,7 @@
    "source": [
     "author = \"OpenBB\"\n",
     "report_title = f\"Economy report\"\n",
-    "user_time_zone = pytz.timezone(helper.get_user_timezone())\n",
+    "user_time_zone = pytz.timezone(get_user_timezone())\n",
     "report_date = pd.Timestamp.now(tz=user_time_zone).strftime(\"%d %B, %Y\")\n",
     "report_time = pd.Timestamp.now(tz=user_time_zone).strftime(\"%H:%M\")\n",
     "report_timezone = pd.Timestamp.now(tz=user_time_zone).tzinfo\n",
@@ -694,7 +694,7 @@
  "metadata": {
   "celltoolbar": "Tags",
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.10.4 ('obb')",
    "language": "python",
    "name": "python3"
   },
@@ -708,11 +708,11 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.6"
+   "version": "3.10.4"
   },
   "vscode": {
    "interpreter": {
-    "hash": "e784374cc4f9a84ed6af2983247aaf351b155c82bcc2188134a65f3c99031000"
+    "hash": "1a8cc6b6e60740679b24fc1ea93bdeb94d949a22102a80c99b7fd3f0d572afd6"
    }
   }
  },

tests/openbb_terminal/economy/test_economy_controller.py

@@ -872,6 +872,7 @@ def test_call_fred_query(mocker):
                 limit=100,
                 raw=False,
                 export="",
+                get_data=True,
             ),
         ),
         (
@@ -885,6 +886,7 @@ def test_call_fred_query(mocker):
                 limit=100,
                 raw=True,
                 export="csv",
+                get_data=True,
             ),
         ),
         (
@@ -898,6 +900,7 @@ def test_call_fred_query(mocker):
                 limit=100,
                 raw=False,
                 export="csv",
+                get_data=True,
             ),
         ),
     ],
@@ -922,7 +925,7 @@ def test_call_fred_params(mocked_func, other_args, called_args, called_kwargs, m
         new=False,
     )
 
-    mock = mocker.Mock()
+    mock = mocker.Mock(return_value=(MOCK_FRED_AGG, MOCK_DETAIL))
     mocker.patch(
         target=f"{path_controller}.{mocked_func}",
         new=mock,