Merge pull request #560 from quantopian/calendar-clean

More flexible calendar implementation

tests/test_finance.py

@@ -526,3 +526,11 @@ class TradingEnvironmentTestCase(TestCase):
         self.assertTrue(all(today == minutes[:31]))
         self.assertTrue(all(friday == minutes[31:421]))
         self.assertTrue(all(thursday == minutes[421:]))
+
+    def test_max_date(self):
+        max_date = datetime(2008, 8, 1, tzinfo=pytz.utc)
+        env = TradingEnvironment(max_date=max_date)
+
+        self.assertLessEqual(env.last_trading_day, max_date)
+        self.assertLessEqual(env.treasury_curves.index[-1],
+                             max_date)

tests/test_tradingcalendar.py

@@ -26,7 +26,6 @@ from nose.tools import nottest
 
 class TestTradingCalendar(TestCase):
 
-    @nottest
     def test_calendar_vs_environment(self):
         """
         test_calendar_vs_environment checks whether the

zipline/data/loader.py

@@ -31,10 +31,8 @@ from six import iteritems
 from . import benchmarks
 from . benchmarks import get_benchmark_returns
 
-from zipline.utils.tradingcalendar import (
-    trading_day,
-    trading_days
-)
+from zipline.utils.tradingcalendar import trading_day as trading_day_nyse
+from zipline.utils.tradingcalendar import trading_days as trading_days_nyse
 
 logger = logbook.Logger('Loader')
 
@@ -154,7 +152,8 @@ def get_benchmark_filename(symbol):
     return "%s_benchmark.csv" % symbol
 
 
-def load_market_data(bm_symbol='^GSPC'):
+def load_market_data(trading_day=trading_day_nyse,
+                     trading_days=trading_days_nyse, bm_symbol='^GSPC'):
     bm_filepath = get_data_filepath(get_benchmark_filename(bm_symbol))
     try:
         saved_benchmarks = pd.Series.from_csv(bm_filepath)

zipline/finance/trading.py

@@ -23,7 +23,6 @@ import numpy as np
 
 from zipline.data.loader import load_market_data
 from zipline.utils import tradingcalendar
-from zipline.utils.tradingcalendar import get_early_closes
 
 
 log = logbook.Logger('Trading')
@@ -93,24 +92,7 @@ class TradingEnvironment(object):
         max_date=None,
         env_trading_calendar=tradingcalendar
     ):
-        self.prev_environment = self
-        self.bm_symbol = bm_symbol
-        if not load:
-            load = load_market_data
-
-        self.benchmark_returns, treasury_curves_map = \
-            load(self.bm_symbol)
-
-        self.treasury_curves = pd.DataFrame(treasury_curves_map).T
-        if max_date:
-            tr_c = self.treasury_curves
-            # Mask the treasury curvers down to the current date.
-            # In the case of live trading, the last date in the treasury
-            # curves would be the day before the date considered to be
-            # 'today'.
-            self.treasury_curves = tr_c[tr_c.index <= max_date]
-
-        self.exchange_tz = exchange_tz
+        self.trading_day = env_trading_calendar.trading_day.copy()
 
         # `tc_td` is short for "trading calendar trading days"
         tc_td = env_trading_calendar.trading_days
@@ -123,12 +105,31 @@ class TradingEnvironment(object):
         self.first_trading_day = self.trading_days[0]
         self.last_trading_day = self.trading_days[-1]
 
-        self.early_closes = get_early_closes(self.first_trading_day,
-                                             self.last_trading_day)
+        self.early_closes = env_trading_calendar.get_early_closes(
+            self.first_trading_day, self.last_trading_day)
 
         self.open_and_closes = env_trading_calendar.open_and_closes.loc[
             self.trading_days]
 
+        self.prev_environment = self
+        self.bm_symbol = bm_symbol
+        if not load:
+            load = load_market_data
+
+        self.benchmark_returns, treasury_curves_map = \
+            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.
+            # In the case of live trading, the last date in the treasury
+            # curves would be the day before the date considered to be
+            # 'today'.
+            self.treasury_curves = tr_c[tr_c.index <= max_date]
+
+        self.exchange_tz = exchange_tz
+
     def __enter__(self, *args, **kwargs):
         global environment
         self.prev_environment = environment