diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 2f1b2d0d..f3b2ffe2 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -123,7 +123,11 @@ class TestEventWindow(TestCase): # Record the length of the window after each event. lengths.append(len(window.ticks)) - assert lengths == [1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3] + # The window stretches out during the weekend because we wait + # to drop events until the weekend ends. The last window is + # briefly longer because it doesn't complete a full day. The + # window then shrinks once the day completes + assert lengths == [1, 2, 3, 3, 3, 4, 5, 5, 5, 3, 4, 3] assert window.added == events assert window.removed == events[:-3] @@ -142,7 +146,7 @@ class TestEventWindow(TestCase): # Record the length of the window after each event. lengths.append(len(window.ticks)) - assert lengths == [1, 2, 2, 2, 2] + assert lengths == [1, 2, 3, 3, 2] assert window.added == events assert window.removed == events[:-2] diff --git a/zipline/transforms/utils.py b/zipline/transforms/utils.py index ea8b7612..4ee0a999 100644 --- a/zipline/transforms/utils.py +++ b/zipline/transforms/utils.py @@ -239,6 +239,9 @@ class EventWindow(object): # adding new ticks. self.handle_add(event) + if self.market_aware: + self.add_new_holidays(event.dt) + # Clear out any expired events. drop_condition changes depending # on whether or not we are running in market_aware mode. # @@ -254,13 +257,33 @@ class EventWindow(object): # behavior for removing ticks. self.handle_remove(popped) - def out_of_market_window(self, oldest, newest): - # Find number of unique days in window - # Note that this assumes that each day we received an - # event is a trading day. - unique_dts = set([event.dt.date() for event in self.ticks]) + def add_new_holidays(self, newest): + # Add to our tracked window any untracked holidays that are + # older than our newest event. (newest should always be + # self.ticks[-1]) + while len(self.all_holidays) > 0 and self.all_holidays[0] <= newest: + self.cur_holidays.append(self.all_holidays.popleft()) - return len(unique_dts) > self.window_length + def drop_old_holidays(self, oldest): + # Drop from our tracked window any holidays that are older + # than our oldest tracked event. (oldest should always + # be self.ticks[0]) + while len(self.cur_holidays) > 0 and self.cur_holidays[0] < oldest: + self.cur_holidays.popleft() + + def out_of_market_window(self, oldest, newest): + self.drop_old_holidays(oldest) + calendar_dates_between = (newest.date() - oldest.date()).days + holidays_between = len(self.cur_holidays) + trading_days_between = calendar_dates_between - holidays_between + + # "Put back" a day if oldest is earlier in its day than newest, + # reflecting the fact that we haven't yet completed the last + # day in the window. + if oldest.time() > newest.time(): + trading_days_between -= 1 + + return trading_days_between >= self.window_length def out_of_delta(self, oldest, newest): return (newest - oldest) >= self.delta