From b4395802424319cfeceae7435a976ee80a0d35f1 Mon Sep 17 00:00:00 2001
From: montezdesousa <79287829+montezdesousa@users.noreply.github.com>
Date: Fri, 22 Jul 2022 00:56:53 +0100
Subject: [PATCH] Fix ycrv to allow raw display and export when source is FRED
 (#2127)

* allow raw display and export

* change title to investpy standard

* limit fred to united states
---
 openbb_terminal/economy/economy_controller.py | 16 +++++++--
 openbb_terminal/economy/fred_model.py         |  6 +++-
 openbb_terminal/economy/fred_view.py          | 34 +++++++++++++++++--
 3 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/openbb_terminal/economy/economy_controller.py b/openbb_terminal/economy/economy_controller.py
index 2a146242886..dfc8d099f29 100644
--- a/openbb_terminal/economy/economy_controller.py
+++ b/openbb_terminal/economy/economy_controller.py
@@ -1008,11 +1008,21 @@ class EconomyController(BaseController):
             if isinstance(ns_parser.country, list):
                 ns_parser.country = " ".join(ns_parser.country)
 
-            investingcom_model.check_correct_country(ns_parser.country)
-
             if ns_parser.source == "FRED":
-                fred_view.display_yield_curve(ns_parser.date)
+
+                if ns_parser.country == "united states":
+                    fred_view.display_yield_curve(
+                        ns_parser.date,
+                        raw=ns_parser.raw,
+                        export=ns_parser.export,
+                    )
+                else:
+                    console.print("Source FRED is only available for united states.\n")
+
             elif ns_parser.source == "investpy":
+
+                investingcom_model.check_correct_country(ns_parser.country)
+
                 investingcom_view.display_yieldcurve(
                     country=ns_parser.country,
                     raw=ns_parser.raw,
diff --git a/openbb_terminal/economy/fred_model.py b/openbb_terminal/economy/fred_model.py
index 97abf261bf3..afabc71b9cb 100644
--- a/openbb_terminal/economy/fred_model.py
+++ b/openbb_terminal/economy/fred_model.py
@@ -257,5 +257,9 @@ def get_yield_curve(date: Optional[datetime]) -> Tuple[pd.DataFrame, str]:
             return pd.DataFrame(), date.strftime("%Y-%m-%d")
         rates = pd.DataFrame(series.values.T, columns=["Rate"])
 
-    rates["Maturity"] = [1 / 12, 1 / 4, 1 / 2, 1, 2, 3, 5, 7, 10, 20, 30]
+    rates.insert(
+        0,
+        "Maturity",
+        ["1M", "3M", "6M", "1Y", "2Y", "3Y", "5Y", "7Y", "10Y", "20Y", "30Y"],
+    )
     return rates, date_of_yield
diff --git a/openbb_terminal/economy/fred_view.py b/openbb_terminal/economy/fred_view.py
index 761ef7eb6d0..b8b94d75e93 100644
--- a/openbb_terminal/economy/fred_view.py
+++ b/openbb_terminal/economy/fred_view.py
@@ -180,7 +180,12 @@ def display_fred_series(
 
 @log_start_end(log=logger)
 @check_api_key(["API_FRED_KEY"])
-def display_yield_curve(date: datetime, external_axes: Optional[List[plt.Axes]] = None):
+def display_yield_curve(
+    date: datetime,
+    external_axes: Optional[List[plt.Axes]] = None,
+    raw: bool = False,
+    export: str = "",
+):
     """Display yield curve based on US Treasury rates for a specified date.
 
     Parameters
@@ -203,10 +208,35 @@ def display_yield_curve(date: datetime, external_axes: Optional[List[plt.Axes]]
     else:
         return
 
-    ax.plot(rates.Maturity, rates.Rate, "-o")
+    tenors = []
+    for i, row in rates.iterrows():
+        t = row["Maturity"][-3:].strip()
+        rates.at[i, "Maturity"] = t
+        if t[-1] == "M":
+            tenors.append(int(t[:-1]) / 12)
+        elif t[-1] == "Y":
+            tenors.append(int(t[:-1]))
+
+    ax.plot(tenors, rates.Rate, "-o")
     ax.set_xlabel("Maturity")
     ax.set_ylabel("Rate (%)")
     theme.style_primary_axis(ax)
     if external_axes is None:
         ax.set_title(f"US Yield Curve for {date_of_yield.strftime('%Y-%m-%d')} ")
         theme.visualize_output()
+
+    if raw:
+        print_rich_table(
+            rates,
+            headers=list(rates.columns),
+            show_index=False,
+            title=f"United States Yield Curve for {date_of_yield.strftime('%Y-%m-%d')}",
+            floatfmt=".3f",
+        )
+
+    export_data(
+        export,
+        os.path.dirname(os.path.abspath(__file__)),
+        "ycrv",
+        rates,
+    )