diff --git a/zipline/transforms/utils.py b/zipline/transforms/utils.py index 0891747f..abf2223a 100644 --- a/zipline/transforms/utils.py +++ b/zipline/transforms/utils.py @@ -29,7 +29,7 @@ from numbers import Integral import pandas as pd from zipline.protocol import Event -from zipline.utils.tradingcalendar import non_trading_days +from zipline.utils import tradingcalendar from zipline.gens.utils import assert_sort_unframe_protocol, hash_args log = logbook.Logger('Transform') @@ -209,9 +209,6 @@ class EventWindow(object): if self.market_aware: assert self.window_length and self.delta is None, \ "Market-aware mode only works with full-day windows." - self.all_holidays = deque(non_trading_days) - self.cur_holidays = deque() - # Non-market-aware mode requires a timedelta. else: assert self.delta and not self.window_length, \ @@ -244,9 +241,6 @@ 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. # @@ -262,25 +256,11 @@ class EventWindow(object): # behavior for removing ticks. self.handle_remove(popped) - 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()) - - 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 + oldest_index = tradingcalendar.trading_days.searchsorted(oldest) + newest_index = tradingcalendar.trading_days.searchsorted(newest) + + trading_days_between = newest_index - oldest_index # "Put back" a day if oldest is earlier in its day than newest, # reflecting the fact that we haven't yet completed the last diff --git a/zipline/utils/tradingcalendar.py b/zipline/utils/tradingcalendar.py index 7de9e621..78dd1bf8 100644 --- a/zipline/utils/tradingcalendar.py +++ b/zipline/utils/tradingcalendar.py @@ -14,6 +14,7 @@ # limitations under the License. +import pandas as pd import pytz from datetime import datetime @@ -176,6 +177,15 @@ def get_non_trading_days(start, end): # - President Gerald R. Ford - Jan 2, 2007 non_trading_days.append(datetime(2007, 1, 2, tzinfo=pytz.utc)) - return sorted(non_trading_days) + return pd.DatetimeIndex(sorted(non_trading_days)) -non_trading_days = get_non_trading_days(start, end) + +def get_trading_days(start, end): + business_days = pd.DatetimeIndex(start=start, end=end, + freq=pd.datetools.BDay()) + + non_trading_days = get_non_trading_days(start, end) + + return business_days - non_trading_days + +trading_days = get_trading_days(start, end)