From a25590b0a15f7b659c31a120697f26e9f65e9c28 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Tue, 1 Jan 2013 13:01:49 -0500 Subject: [PATCH] Exposes the list of trading days contained in a trading environment. Previously, the list was generated, but only used to calculate the number of days in the environment. With exposing this list, working towards a path where the simulation uses the trading days to determine when to handle market closes. --- tests/test_finance.py | 42 ++++++++++++++++++++++++++++++++++++++ zipline/finance/trading.py | 16 +++++++-------- 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/tests/test_finance.py b/tests/test_finance.py index 22dbc23f..af4fc91d 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -21,6 +21,8 @@ import pytz from unittest import TestCase from datetime import datetime, timedelta +import numpy as np + from nose.tools import timed import zipline.utils.factory as factory @@ -117,6 +119,46 @@ class FinanceTestCase(TestCase): self.assertTrue(env.last_close.month == 12) self.assertTrue(env.last_close.day == 31) + @timed(DEFAULT_TIMEOUT) + def test_trading_environment_days_in_period(self): + + benchmark_returns, treasury_curves = \ + factory.load_market_data() + + # January 2008 + # Su Mo Tu We Th Fr Sa + # 1 2 3 4 5 + # 6 7 8 9 10 11 12 + # 13 14 15 16 17 18 19 + # 20 21 22 23 24 25 26 + # 27 28 29 30 31 + + env = TradingEnvironment( + benchmark_returns, + treasury_curves, + period_start=datetime(2007, 12, 31, tzinfo=pytz.utc), + period_end=datetime(2008, 1, 7, tzinfo=pytz.utc), + capital_base=100000, + ) + + expected_trading_days = ( + datetime(2007, 12, 31, tzinfo=pytz.utc), + # Skip new years + #holidays taken from: http://www.nyse.com/press/1191407641943.html + datetime(2008, 1, 2, tzinfo=pytz.utc), + datetime(2008, 1, 3, tzinfo=pytz.utc), + datetime(2008, 1, 4, tzinfo=pytz.utc), + # Skip Saturday + # Skip Sunday + datetime(2008, 1, 7, tzinfo=pytz.utc) + ) + + num_expected_trading_days = 5 + + self.assertEquals(num_expected_trading_days, env.days_in_period) + np.testing.assert_array_equal(expected_trading_days, + env.period_trading_days) + @timed(EXTENDED_TIMEOUT) def test_full_zipline(self): #provide enough trades to ensure all orders are filled. diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 0f60b742..8c979562 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -73,7 +73,7 @@ class TradingEnvironment(object): self.period_start = period_start self.period_end = period_end self.capital_base = capital_base - self.period_trading_days = None + self._period_trading_days = None assert self.period_start <= self.period_end, \ "Period start falls after period end." @@ -170,19 +170,19 @@ class TradingEnvironment(object): ) @property - def days_in_period(self): - """return the number of trading days within the period [start, end)""" - assert self.period_start is not None - assert self.period_end is not None - - if self.period_trading_days is None: - self.period_trading_days = [] + def period_trading_days(self): + if self._period_trading_days is None: + self._period_trading_days = [] for date in self.trading_day_map.iterkeys(): if date > self.period_end: break if date >= self.period_start: self.period_trading_days.append(date) + return self._period_trading_days + @property + def days_in_period(self): + """return the number of trading days within the period [start, end)""" return len(self.period_trading_days) def is_market_hours(self, test_date):