diff --git a/tests/test_finance.py b/tests/test_finance.py index b511347f..87ac9858 100644 --- a/tests/test_finance.py +++ b/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) diff --git a/tests/test_tradingcalendar.py b/tests/test_tradingcalendar.py index c5a43227..4ea8e498 100644 --- a/tests/test_tradingcalendar.py +++ b/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 diff --git a/zipline/data/loader.py b/zipline/data/loader.py index 3884cb7c..9b6d7d21 100644 --- a/zipline/data/loader.py +++ b/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) diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 44b516b6..6c3a1e47 100644 --- a/zipline/finance/trading.py +++ b/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