diff --git a/tests/test_finance.py b/tests/test_finance.py index 4197de07..ea3a3ce1 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -25,6 +25,7 @@ from unittest import TestCase from datetime import datetime, timedelta import numpy as np +import pandas as pd from nose.tools import timed @@ -546,8 +547,16 @@ class TradingEnvironmentTestCase(TestCase): self.assertTrue(all(friday == minutes[31:421])) self.assertTrue(all(thursday == minutes[421:])) + def test_min_date(self): + min_date = pd.Timestamp('2016-03-04', tz='UTC') + env = TradingEnvironment(min_date=min_date) + + self.assertGreaterEqual(env.first_trading_day, min_date) + self.assertGreaterEqual(env.treasury_curves.index[0], + min_date) + def test_max_date(self): - max_date = datetime(2008, 8, 1, tzinfo=pytz.utc) + max_date = pd.Timestamp('2008-08-01', tz='UTC') env = TradingEnvironment(max_date=max_date) self.assertLessEqual(env.last_trading_day, max_date) diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index d914c611..362d5534 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -68,6 +68,7 @@ class TradingEnvironment(object): load=None, bm_symbol='^GSPC', exchange_tz="US/Eastern", + min_date=None, max_date=None, env_trading_calendar=tradingcalendar, asset_db_path=':memory:' @@ -82,10 +83,7 @@ class TradingEnvironment(object): # `tc_td` is short for "trading calendar trading days" tc_td = env_trading_calendar.trading_days - if max_date: - self.trading_days = tc_td[tc_td <= max_date].copy() - else: - self.trading_days = tc_td.copy() + self.trading_days = tc_td[tc_td.slice_indexer(min_date, max_date)] self.first_trading_day = self.trading_days[0] self.last_trading_day = self.trading_days[-1]