From 0fa44471bed0f26b9ac27b7ec591e767691dee1a Mon Sep 17 00:00:00 2001
From: Eddie Hebert <ehebert@quantopian.com>
Date: Mon, 20 Apr 2015 09:49:01 -0400
Subject: [PATCH] MAINT: Change expected type of treasury curves from load to
 DataFrame.

Instead of converting the curves back and forth from dictionaries to
DataFrame and back, use the DataFrame format when passing to
environment.
---
 zipline/data/loader.py     | 12 +-----------
 zipline/finance/trading.py |  8 ++++++--
 2 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/zipline/data/loader.py b/zipline/data/loader.py
index 347dcff4..76e335f9 100644
--- a/zipline/data/loader.py
+++ b/zipline/data/loader.py
@@ -229,17 +229,7 @@ def load_market_data(trading_day=trading_day_nyse,
     else:
         treasury_curves = saved_curves.tz_localize('UTC')
 
-    tr_curves = {}
-    for tr_dt, curve in treasury_curves.T.iteritems():
-        # tr_dt = tr_dt.replace(hour=0, minute=0, second=0, microsecond=0,
-        #                       tzinfo=pytz.utc)
-        tr_curves[tr_dt] = curve.to_dict()
-
-    tr_curves = OrderedDict(sorted(
-        ((dt, c) for dt, c in iteritems(tr_curves)),
-        key=lambda t: t[0]))
-
-    return benchmark_returns, tr_curves
+    return benchmark_returns, treasury_curves
 
 
 def _load_raw_yahoo_data(indexes=None, stocks=None, start=None, end=None):
diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py
index 6c3a1e47..3ac5c7cb 100644
--- a/zipline/finance/trading.py
+++ b/zipline/finance/trading.py
@@ -92,6 +92,11 @@ class TradingEnvironment(object):
         max_date=None,
         env_trading_calendar=tradingcalendar
     ):
+        """
+        @load is function that returns benchmark_returns and treasury_curves
+        The treasury_curves are expected to be a DataFrame with an index of
+        dates and columns of the curve names, e.g. '10year', '1month', etc.
+        """
         self.trading_day = env_trading_calendar.trading_day.copy()
 
         # `tc_td` is short for "trading calendar trading days"
@@ -116,10 +121,9 @@ class TradingEnvironment(object):
         if not load:
             load = load_market_data
 
-        self.benchmark_returns, treasury_curves_map = \
+        self.benchmark_returns, self.treasury_curves = \
             load(self.trading_day, self.trading_days, self.bm_symbol)
 
-        self.treasury_curves = pd.DataFrame(treasury_curves_map).T
         if max_date:
             tr_c = self.treasury_curves
             # Mask the treasury curves down to the current date.