From 5786c7de6e68e9601344fc4eafb93a80bac0ae32 Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Wed, 25 Nov 2015 11:54:56 +0100
Subject: [PATCH] ENH Add treasury loader with various maturities.

---
 pyfolio/utils.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/pyfolio/utils.py b/pyfolio/utils.py
index 030b56ab..a776e679 100644
--- a/pyfolio/utils.py
+++ b/pyfolio/utils.py
@@ -285,6 +285,38 @@ def load_portfolio_risk_factors(filepath_prefix=None, start=None, end=None):
     return five_factors.loc[start:end]
 
 
+def get_treasury_yield(start=None, end=None, period='3MO'):
+    """Load treasury yields from FRED.
+
+    Parameters
+    ----------
+    start : date, optional
+        Earliest date to fetch data for.
+        Defaults to earliest date available.
+    end : date, optional
+        Latest date to fetch data for.
+        Defaults to latest date available.
+    period : {'1MO', '3MO', '6MO', 1', '5', '10'}, optional
+        Which maturity to use.
+
+    Returns
+    -------
+    pd.Series
+        Annual treasury yield for every day.
+    """
+    if start is None:
+        start = '1/1/1970'
+    if end is None:
+        end = pd.Timestamp(datetime.today()).normalize() - BDay()
+
+    treasury = web.DataReader("DGS3{}".format(period), "fred",
+                              start, end)
+
+    treasury = treasury.ffill()
+
+    return treasury
+
+
 def extract_rets_pos_txn_from_zipline(backtest):
     """Extract returns, positions, transactions and leverage from the
     backtest data structure returned by zipline.TradingAlgorithm.run().