From 69b87b859bd50e224692388976d863d329fa62bf Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Thu, 31 Jan 2013 13:06:15 -0500 Subject: [PATCH] Fixes trading day generation logic in test source. This uses more of @fawce's original logic from record_var branch. TODO: Add unit test to ensure that this function generates only market days. --- zipline/sources/test_source.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/zipline/sources/test_source.py b/zipline/sources/test_source.py index f589392a..a3720f8d 100644 --- a/zipline/sources/test_source.py +++ b/zipline/sources/test_source.py @@ -34,6 +34,7 @@ def date_gen(start=datetime(2006, 6, 6, 12, tzinfo=pytz.utc), """ Utility to generate a stream of dates. """ + one_day = timedelta(days=1) cur = start # yield count trade events, all on trading days, and @@ -50,9 +51,10 @@ def date_gen(start=datetime(2006, 6, 6, 12, tzinfo=pytz.utc), cur = cur + delta cur_midnight = cur.replace(hour=0, minute=0, second=0) # skip over any non-trading days - if cur_midnight not in trading_days: - next_day_index = trading_days.searchsorted(cur_midnight) - cur_midnight = trading_days[next_day_index] + while cur_midnight not in trading_days: + cur = cur + one_day + cur_midnight = cur.replace(hour=0, minute=0, second=0) + cur = cur.replace(day=cur_midnight.day) def mock_prices(count):