From f4c48dc8d448bd6fd989ddab8d550a36816c3fb6 Mon Sep 17 00:00:00 2001
From: Colin Delahunty <72827203+colin99d@users.noreply.github.com>
Date: Mon, 3 Apr 2023 05:07:36 -0700
Subject: [PATCH] Improves `economy/fred` command (#4634)

* Fixed

* Started the switch

---------

Co-authored-by: James Maslek <jmaslek11@gmail.com>
---
 openbb_terminal/economy/economy_controller.py | 13 ++++--
 openbb_terminal/economy/fred_view.py          | 42 ++++++++++++++-----
 .../economy/test_economy_controller.py        |  4 +-
 3 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/openbb_terminal/economy/economy_controller.py b/openbb_terminal/economy/economy_controller.py
index 45ac42386a7..0fcd3e795cc 100644
--- a/openbb_terminal/economy/economy_controller.py
+++ b/openbb_terminal/economy/economy_controller.py
@@ -1389,8 +1389,13 @@ class EconomyController(BaseController):
                 query = " ".join(ns_parser.query)
                 df_search = fred_model.get_series_notes(search_query=query)
 
-                if not df_search.empty:
-                    fred_view.notes(search_query=query, limit=ns_parser.limit)
+                if isinstance(df_search, pd.DataFrame) and not df_search.empty:
+                    fred_view.notes(
+                        search_query=query,
+                        limit=ns_parser.limit,
+                        export=ns_parser.export,
+                        sheet_name=ns_parser.sheet_name,
+                    )
 
                     self.fred_query = df_search["id"].head(ns_parser.limit)
                     self.update_runtime_choices()
@@ -1431,7 +1436,7 @@ class EconomyController(BaseController):
                     get_data=True,
                 )
 
-                if not df.empty:
+                if isinstance(df, pd.DataFrame) and not df.empty:
                     for series_id, data in detail.items():
                         self.FRED_TITLES[
                             series_id
@@ -1448,7 +1453,7 @@ class EconomyController(BaseController):
                     if get_current_user().preferences.ENABLE_EXIT_AUTO_HELP:
                         self.print_help()
 
-                else:
+                elif not ns_parser.export and not ns_parser.raw:
                     console.print("[red]No data found for the given Series ID[/red]")
 
             elif not parameters and ns_parser.raw:
diff --git a/openbb_terminal/economy/fred_view.py b/openbb_terminal/economy/fred_view.py
index d1b5d85f3de..9cddc52fa0f 100644
--- a/openbb_terminal/economy/fred_view.py
+++ b/openbb_terminal/economy/fred_view.py
@@ -40,7 +40,12 @@ 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):
+def notes(
+    search_query: str,
+    limit: int = 10,
+    export: str = "",
+    sheet_name: Optional[str] = None,
+):
     """Display series notes. [Source: FRED]
 
     Parameters
@@ -49,6 +54,10 @@ def notes(search_query: str, limit: int = 10):
         Text query to search on fred series notes database
     limit : int
         Maximum number of series notes to display
+    export : str
+        Export data to csv,json,xlsx or png,jpg,pdf,svg file
+    sheet_name : Optional[str]
+        The name of the sheet
     """
     df_search = fred_model.get_series_notes(search_query)
 
@@ -61,6 +70,14 @@ def notes(search_query: str, limit: int = 10):
         show_index=False,
         headers=["Series ID", "Title", "Description"],
         limit=limit,
+        export=bool(export),
+    )
+    export_data(
+        export,
+        os.path.dirname(os.path.abspath(__file__)),
+        "fred",
+        df_search,
+        sheet_name,
     )
 
 
@@ -93,6 +110,8 @@ def display_fred_series(
         Output only raw data
     export : str
         Export data to csv,json,xlsx or png,jpg,pdf,svg file
+    sheet_name : Optional[str]
+        The name of the sheet
     external_axes : bool, optional
         Whether to return the figure object or not, by default False
     """
@@ -132,6 +151,17 @@ def display_fred_series(
     )
     data.index = [x.strftime("%Y-%m-%d") for x in data.index]
 
+    if export:
+        export_data(
+            export,
+            os.path.dirname(os.path.abspath(__file__)),
+            "fred",
+            data,
+            sheet_name,
+            fig,
+        )
+        return None, None
+
     if raw:
         data = data.sort_index(ascending=False)
         print_rich_table(
@@ -142,15 +172,7 @@ def display_fred_series(
             export=bool(export),
             limit=limit,
         )
-
-    export_data(
-        export,
-        os.path.dirname(os.path.abspath(__file__)),
-        "fred",
-        data,
-        sheet_name,
-        fig,
-    )
+        return None, None
 
     if get_data:
         fig.show(external=external_axes)
diff --git a/tests/openbb_terminal/economy/test_economy_controller.py b/tests/openbb_terminal/economy/test_economy_controller.py
index 46f19263c86..52e62de5520 100644
--- a/tests/openbb_terminal/economy/test_economy_controller.py
+++ b/tests/openbb_terminal/economy/test_economy_controller.py
@@ -1040,7 +1040,9 @@ def test_call_fred_query(mocker):
     controller = economy_controller.EconomyController(queue=None)
     controller.choices = {}
     controller.call_fred(["--query", "MOCK_QUERY", "--limit", "1"])
-    mock.assert_called_once_with(**dict(search_query="MOCK_QUERY", limit=1))
+    mock.assert_called_once_with(
+        **dict(search_query="MOCK_QUERY", limit=1, export="", sheet_name=None)
+    )
 
 
 @pytest.mark.vcr(record_mode="none")