Adding a Yield Curve model

gamestonk_terminal/fundamental_analysis/yield_curve_model.py

@@ -0,0 +1,95 @@
+""" Fundamental Analysis Yield Curve Model """
+__docformat__ = "numpy"
+
+from datetime import datetime
+from bs4 import BeautifulSoup
+import pandas as pd
+from pandas.core.frame import DataFrame
+import requests
+
+from gamestonk_terminal.helper_funcs import (
+    get_user_agent,
+)
+
+
+def get_yield_curve(start: datetime, end: datetime) -> DataFrame:
+    """Returns a US Treasury Yield Curve from start date to end date
+
+    Parameters
+    ----------
+    start : datetime
+        Start date
+    end : datetime
+        End date
+
+    Returns
+    -------
+    DataFrame
+        US Treasury Yield Curve data frame with the following columns
+        Date,1 mo,2 mo,3 mo,6 mo,1 yr,2 yr,3 yr,5 yr,7 yr,10 yr,20 yr,30 yr
+    """
+
+    if start.year == end.year:
+        df_yield_curve = get_yield_curve_year(start.year)
+        return df_yield_curve[start:end]
+
+    df_yield_curve = get_yield_curve_year(start.year)
+
+    a_year = start.year + 1
+
+    while a_year <= end.year:
+        df_temp_curve = get_yield_curve_year(a_year)
+        df_yield_curve = df_yield_curve.append(df_temp_curve)
+        a_year += 1
+
+    return df_yield_curve[start:end]
+
+
+def get_yield_curve_year(year: str) -> DataFrame:
+    """Returns a US Treasury Yield Curve for a given year
+
+    Parameters
+    ----------
+    year : str
+        Yield curve year
+
+    Returns
+    -------
+    DataFrame
+        US Treasury Yield Curve data frame with the following columns
+        Date,1 mo,2 mo,3 mo,6 mo,1 yr,2 yr,3 yr,5 yr,7 yr,10 yr,20 yr,30 yr
+    """
+
+    # pylint: disable=line-too-long
+    yield_curve_url = "https://www.treasury.gov/resource-center/data-chart-center/interest-rates/Pages/TextView.aspx?data=yieldYear&year={}"  # noqa: E501
+    text_soup_yield_curve = BeautifulSoup(
+        requests.get(
+            yield_curve_url.format(year),
+            headers={"User-Agent": get_user_agent()},
+        ).text,
+        "lxml",
+    )
+
+    yield_table = text_soup_yield_curve.find_all("table", {"class": "t-chart"})[0]
+
+    a_yield_table_header = list()
+    for yield_table_header_col in yield_table.find("tr", {"class": None}).find_all(
+        "th"
+    ):
+        a_yield_table_header.append(yield_table_header_col.text)
+
+    df_yield_curve = pd.DataFrame(columns=a_yield_table_header)
+
+    for yield_table_row in yield_table.find_all("tr", {"class": ["oddrow", "evenrow"]}):
+        a_yield_row = list()
+        for idx, yield_table_col in enumerate(yield_table_row.find_all("td")):
+            if idx == 0:
+                a_yield_row.append(datetime.strptime(yield_table_col.text, "%m/%d/%y"))
+            else:
+                a_yield_row.append(float(yield_table_col.text))
+
+        df_yield_curve.loc[len(df_yield_curve)] = a_yield_row
+
+    df_yield_curve = df_yield_curve.set_index(["Date"])
+
+    return df_yield_curve

gamestonk_terminal/test_helper.py

@@ -6,7 +6,7 @@ from typing import Callable
 import yaml
 
 
-def parametrize_from_file(test_namespace: str, parameter_file: str) -> Callable:
+def parameterize_from_file(test_namespace: str, parameter_file: str) -> Callable:
     """A test helper function returns a wrapper function that loads sets of parameters from a YAML and
     attaches test scenarios as metadata
 

tests/test_fa_market_watch_api_unit.py

@@ -4,7 +4,7 @@ from unittest import mock
 
 # pylint: disable=unused-import
 from gamestonk_terminal.test_helper import (  # noqa: F401
-    parametrize_from_file,
+    parameterize_from_file,
     pytest_generate_tests,
 )
 
@@ -17,7 +17,7 @@ assertions = unittest.TestCase("__init__")
 
 class TestFaMarketWatchApiUnit:
     @mock.patch("gamestonk_terminal.fundamental_analysis.market_watch_api.requests")
-    @parametrize_from_file(
+    @parameterize_from_file(
         "test_prepare_df_financials",
         "../tests/data/fa_market_watch_api.yaml",
     )

tests/test_fa_yield_curve_model.py

@@ -0,0 +1,70 @@
+""" fundamental_analysis/yield_curve_model.py tests """
+from datetime import datetime
+from io import StringIO
+import unittest
+from unittest import mock
+import pandas as pd
+
+# pylint: disable=unused-import
+from gamestonk_terminal.test_helper import (  # noqa: F401
+    parameterize_from_file,
+    pytest_generate_tests,
+)
+
+from gamestonk_terminal.fundamental_analysis.yield_curve_model import (
+    get_yield_curve,
+    get_yield_curve_year,
+)
+
+assertions = unittest.TestCase("__init__")
+
+
+class TestFaYieldCurveModel:
+    @mock.patch(
+        "gamestonk_terminal.fundamental_analysis.yield_curve_model.get_yield_curve_year"
+    )
+    @parameterize_from_file(
+        "test_get_yield_curve",
+        "../tests/data/fa_yield_curve_model.yaml",
+    )
+    # pylint: disable=too-many-arguments
+    def test_get_yield_curve(
+        self,
+        mock_get_yield_curve_year,
+        start,
+        end,
+        mock_get_yield_curve_year_rets,
+        expected_result,
+    ):
+
+        rets = list()
+        for a_ret in mock_get_yield_curve_year_rets:
+            rets.append(
+                pd.read_csv(StringIO(a_ret), header=0, index_col=0, parse_dates=True)
+            )
+
+        mock_get_yield_curve_year.side_effect = rets
+        df = get_yield_curve(
+            datetime.strptime(start, "%m/%d/%y"),
+            datetime.strptime(end, "%m/%d/%y"),
+        )
+
+        assertions.assertEqual(
+            df.to_csv().replace("\r\n", "\n"), expected_result.replace("\r\n", "\n")
+        )
+
+    @mock.patch("gamestonk_terminal.fundamental_analysis.yield_curve_model.requests")
+    @parameterize_from_file(
+        "test_get_yield_curve_year",
+        "../tests/data/fa_yield_curve_model.yaml",
+    )
+    def test_get_yield_curve_year(
+        self, mock_request_get, year, mock_yield_curve_page, expected_result
+    ):
+        mock_request_get.get().text = mock_yield_curve_page
+
+        ret = get_yield_curve_year(year)
+
+        assertions.assertEqual(
+            ret.to_csv().replace("\r\n", "\n"), expected_result.replace("\r\n", "\n")
+        )

tests/test_thought_of_the_day.py

@@ -4,7 +4,7 @@ from unittest import mock
 
 # pylint: disable=unused-import
 from gamestonk_terminal.test_helper import (  # noqa: F401
-    parametrize_from_file,
+    parameterize_from_file,
     pytest_generate_tests,
 )
 from gamestonk_terminal import thought_of_the_day
@@ -13,7 +13,7 @@ assertions = unittest.TestCase("__init__")
 
 
 class TestThoughtOfTheDay:
-    @parametrize_from_file(
+    @parameterize_from_file(
         "test_get_urls",
         "../tests/data/thought_of_the_day.yaml",
     )
@@ -23,7 +23,7 @@ class TestThoughtOfTheDay:
         assertions.assertEqual(a_totd.get_urls(), urls)
 
     @mock.patch("gamestonk_terminal.thought_of_the_day.requests")
-    @parametrize_from_file(
+    @parameterize_from_file(
         "test_get_metadata",
         "../tests/data/thought_of_the_day.yaml",
     )