diff --git a/tests/test_transforms.py b/tests/test_transforms.py index e515d725..8e203664 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -46,24 +46,20 @@ class EventWindowTestCase(TestCase): def setUp(self): setup_logger(self) - # Constants calling before open, during the day, and after - # close on a valid trading day. - self.pre_open = datetime(2012, 8, 7, 13, tzinfo = pytz.utc) - self.mid_day = datetime(2012, 8, 7, 15, tzinfo = pytz.utc) - self.post_close = datetime(2012, 8, 7, 22, tzinfo = pytz.utc) + self.monday = datetime(2012, 7, 9, 16, tzinfo=pytz.utc) + self.eleven_normal_days = [self.monday + i*timedelta(days=1) + for i in xrange(11)] - # Constants calling before open, during the day, and after - # close on a saturday. - self.pre_open_saturday = datetime(2012, 8, 11, 13, tzinfo = pytz.utc) - self.mid_day_saturday = datetime(2012, 8, 11, 15, tzinfo = pytz.utc) - self.post_close_saturday = datetime(2012, 8, 11, 22, tzinfo = pytz.utc) + # Modify the end of the period slightly to exercise the + # incomplete day logic. + self.eleven_normal_days[-1] -= timedelta(minutes = 1) + self.eleven_normal_days.append(self.monday+timedelta(days=11,seconds=1)) + + # Second set of dates to test holiday handling. + self.jul4_monday = datetime(2012, 7, 2, 16, tzinfo=pytz.utc) + self.week_of_jul4 = [self.jul4_monday + i*timedelta(days=1) + for i in xrange(5)] - # Constants calling before open, during the day, and after - # close on a holiday. - self.pre_open_holiday = datetime(2012, 12, 25, 13, tzinfo = pytz.utc) - self.mid_day_holiday = datetime(2012, 12, 25, tzinfo = pytz.utc) - self.post_close_holiday = datetime(2012, 12, 25, 22, tzinfo = pytz.utc) - def test_event_window_with_timedelta(self): # Keep all events within a 5 minute window. @@ -96,58 +92,46 @@ class EventWindowTestCase(TestCase): for dropped in window.removed: assert message.dt - dropped.dt >= timedelta(minutes = 5) - def test_market_aware_window(self): + def test_market_aware_window_normal_week(self): window = NoopEventWindow( market_aware = True, delta = None, - days = 1 + days = 3 ) - - dates = ([self.pre_open]*3) - dates += ([self.mid_day]*3) - dates += ([self.post_close]*3) - dates += [self.pre_open + timedelta(days = 1, seconds = 1)] - events = [to_dt(date) for date in dates] - + events = [to_dt(date) for date in self.eleven_normal_days] + lengths = [] # Run the events. for event in events: window.update(event) + # Record the length of the window after each event. + lengths.append(len(window.ticks)) - # We should have removed the pre_open events on the first day. - # The rest should be intact. - + # 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[0:3] - assert list(window.ticks) == events[3:] + assert window.removed == events[:-3] - def test_market_aware_window_weekend(self): + def test_market_aware_window_holiday(self): window = NoopEventWindow( market_aware = True, delta = None, days = 2 ) - dates = [self.pre_open_saturday - timedelta(days = 1, seconds=1)] - dates += [self.mid_day_saturday - timedelta(days = 1, seconds=1)] - dates += [self.post_close_saturday - timedelta(days = 1, seconds=1)] - dates += [self.mid_day_saturday + timedelta(days = 1)] - - events = [to_dt(date) for date in dates] + events = [to_dt(date) for date in self.week_of_jul4] + lengths = [] # Run the events. for event in events: window.update(event) + # Record the length of the window after each event. + lengths.append(len(window.ticks)) - # We shouldn't remove any events. + assert lengths == [1, 2, 3, 3, 2] assert window.added == events - assert window.removed == [] - assert list(window.ticks) == events - - extra = to_dt(self.mid_day_saturday + timedelta(days = 2)) - window.update(extra) - - # We should remove only the first event. - assert window.removed == [events[0]] - assert list(window.ticks) == events[1:] + [extra] + assert window.removed == events[:-2] def tearDown(self): setup_logger(self) diff --git a/zipline/finance/performance.py b/zipline/finance/performance.py index 05a3a875..de7809d2 100644 --- a/zipline/finance/performance.py +++ b/zipline/finance/performance.py @@ -219,7 +219,7 @@ class PerformanceTracker(object): del event['TRANSACTION'] yield event # Cut off the rest of the stream. - yield StopIteration() + raise StopIteration() else: event.perf_message = self.process_event(event) event.portfolio = self.get_portfolio() diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index cec8b448..63d69d0f 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -77,11 +77,11 @@ def sequential_transforms(stream_in, *transforms): """ assert isinstance(transforms, (list, tuple)) - for tnfm in transforms: - tnfm.forward_all = False - tnfm.update_in_place = False - tnfm.append_value = True + for tnfm in transforms: + tnfm.sequential = True + tnfm.merged = False + # Recursively apply all transforms to the stream. stream_out = reduce(lambda stream, tnfm: tnfm.transform(stream), transforms, diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index e10380b6..33827448 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -12,7 +12,7 @@ from numbers import Number from abc import ABCMeta, abstractmethod from zipline import ndict -from zipline.utils.tradingcalendar import trading_days_between +from zipline.utils.tradingcalendar import non_trading_days from zipline.gens.utils import assert_sort_unframe_protocol, \ assert_transform_protocol, hash_args @@ -48,8 +48,11 @@ class StatefulTransform(object): # behavior if we are being fed to merged_transforms. self.passthrough = tnfm_class.__dict__.get('PASSTHROUGH', False) - self.sequential = True - self.merged = False + # Flags specifying how to append the calculated value. + # Merged is the default for ease of testing, but we use sequential + # in production. + self.sequential = False + self.merged = True # Create an instance of our transform class. self.state = tnfm_class(*args, **kwargs) @@ -120,8 +123,9 @@ class StatefulTransform(object): out_message = message out_message[self.namestring] = tnfm_value yield out_message - + log.info('Finished StatefulTransform [%s]' % self.get_hash()) + class EventWindow: """ Abstract base class for transform classes that calculate iterative @@ -153,8 +157,13 @@ class EventWindow: # Market-aware mode only works with full-day windows. if self.market_aware: - assert self.days and not self.delta,\ + assert self.days and self.delta == None,\ "Market-aware mode only works with full-day windows." + self.all_holidays = deque(non_trading_days) + self.cur_holidays = deque() + # Keeping a copy of days as a timedelta makes it easier + # to track holidays. + self.delta = timedelta(days=self.days) # Non-market-aware mode requires a timedelta. else: @@ -188,6 +197,9 @@ class EventWindow: # Subclasses should override handle_add to define behavior for # 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. @@ -196,16 +208,40 @@ class EventWindow: # | | # V V while self.drop_condition(self.ticks[0].dt, self.ticks[-1].dt): - + # popleft removes and returns the oldest tick in self.ticks popped = self.ticks.popleft() # Subclasses should override handle_remove to define # 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): - return trading_days_between(oldest, newest) >= self.days + 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.days def out_of_delta(self, oldest, newest): return (newest - oldest) >= self.delta diff --git a/zipline/utils/tradingcalendar.py b/zipline/utils/tradingcalendar.py index f760e51e..36dd3eea 100644 --- a/zipline/utils/tradingcalendar.py +++ b/zipline/utils/tradingcalendar.py @@ -1,358 +1,114 @@ import pytz -from datetime import datetime, timedelta +from datetime import datetime, timedelta, date from dateutil import rrule from zipline.utils.date_utils import utcnow -def market_opens(start, end, inclusive=False): - """ - Returns all market opens between the start date and the end date. - Must use utc-stamped datetimes. - """ - return opens.between(start, end, inc=inclusive) +start = datetime(2002, 1,1, tzinfo=pytz.utc) +end = utcnow() -def market_closes(start, end, inclusive=False): - """ - Returns all market closes between the start date and the end date. - Must use utc-stamped datetimes. - """ - return closes.between(start, end, inc=inclusive) +non_trading_rules = [] -def trading_days_between(start, end): - """ - Calculate the number of "complete" trading days between two - events. We define this as the number of market opens that - occurred between start and end, with the caveat that we subtract 1 - from this total if end falls on the same day as the last market - open and end occurs earlier in its own day than start. This - reflects the fact that we haven't completed a full day - corresponding to the last market open. - - Examples: - - 1.) - start = Tuesday, Aug 7, 2012, 1:00 pm - end = Wednesday, Aug 8, 2012, 1:30 pm - - There is one market open between these dates, on the morning of - Wednesday the 8th. This falls on the same calendar day as end, - but end is later in the day than start, so we count this as a full - day. The correct output is 1. - - 2.) - start = Tuesday, Aug 7, 2012, 1:30 pm - end = Wednesday, Aug 8, 2012, 1:00 pm - - There is one market open between these dayes, on the morning of - Wednesday the 8th. This falls on the same calendar day as end, - and end is earlier in the day than start, so we do not count this - day as completed. The correct output is 0. - - 3.) - start = Tuesday, Aug 7, 2012, 1:00 pm - end = Saturday, Aug 11, 2012, 1:30 pm - - There are 3 market opens between these dates, occurring on - Wednesday, Thursday, and Friday. The last open is not on - the same day as end, so we simply return 3 - - 4.) - start = Tuesday, Aug 7, 2012, 1:30 pm - end = Monday, Aug, 13, 2012, 1:00 pm - - There are 4 market opens between these dates, occurring on - Wednesday, Thursday, Friday, and the following Monday. The - last open occurs on the same calendar day as end, and end - is earlier in the day than start, so we do not count the - last market day as completed. The correct output is 3 days. - """ - # Calculate the number of opens between the events. - opens = (market_opens(start, end)) - days_between = len(opens) - if days_between == 0: - return days_between - - # If end falls on the same day as an open, subtract 1 from the - # total if end is earlier in its respective day than start. - last_open = opens[-1] - if last_open.date() == end.date() and earlier_in_day(end, start): - days_between -=1 - - return days_between - -def earlier_in_day(d1, d2): - """ - Return true if d1 falls earlier in its own day than d2. - """ - return d1.time() < d2.time() - -WEEKDAYS = [rrule.MO, rrule.TU, rrule.WE, rrule.TH, rrule.FR] - -# Recurrence rule that generates all market opens since Jan 1, 1970. -# This does not exclude holidays. -market_opens_with_holidays = rrule.rrule( - rrule.DAILY, - byweekday=WEEKDAYS, - byhour = 14, - byminute = 30, +weekends = rrule.rrule( + rrule.YEARLY, + byweekday=(rrule.SA, rrule.SU), cache = True, - dtstart=datetime(2000, 1, 1, tzinfo = pytz.utc), - until=datetime(2014 , 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) +non_trading_rules.append(weekends) -# Recurrence rule that generates all market closes since Jan 1, 1970. -# This does not exclude holidays. -market_closes_with_holidays = rrule.rrule( - rrule.DAILY, - byweekday=WEEKDAYS, - byhour = 21, - byminute = 0, - cache = True, - dtstart=datetime(2001, 1, 1, tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) - -# Recurrence rules for excluding the market open/close on new years. -new_years_opens = rrule.rrule( +new_years = rrule.rrule( rrule.MONTHLY, byyearday = 1, - byhour = 14, - byminute = 30, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) -new_years_closes = rrule.rrule( - rrule.MONTHLY, - byyearday = 1, - byhour = 21, - byminute = 0, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) +non_trading_rules.append(new_years) -# Recurrence rules for excluding MLK day. It is always the third -# monday in January. -mlk_opens = rrule.rrule( - rrule.MONTHLY, - bymonth = 1, - byweekday = (rrule.MO(3)), - byhour = 14, - byminute = 30, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) -mlk_closes = rrule.rrule( +mlk_day = rrule.rrule( rrule.MONTHLY, bymonth = 1, byweekday = (rrule.MO(+3)), - byhour = 21, - byminute = 0, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) + dtstart = start, + until = end +) +non_trading_rules.append(mlk_day) -# Recurrence rules for generating the market open/close for -# presidents' day. Presidents' day always occurs on the third monday -# of February. -presidents_day_opens = rrule.rrule( +presidents_day = rrule.rrule( rrule.MONTHLY, bymonth = 2, byweekday = (rrule.MO(3)), - byhour = 14, - byminute = 30, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) -presidents_day_closes = rrule.rrule( - rrule.MONTHLY, - bymonth = 2, - byweekday = (rrule.MO(3)), - byhour = 21, - byminute = 0, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) +non_trading_rules.append(presidents_day) -# Recurrence rules for generating the market open/close for good -# friday. Good friday always falls 2 days before easter, which -# thankfully is a built-in refernce in this module. -good_friday_opens = rrule.rrule( +good_friday = rrule.rrule( rrule.DAILY, byeaster = -2, - byhour = 14, - byminute = 30, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) -good_friday_closes = rrule.rrule( - rrule.DAILY, - byeaster = -2, - byhour = 21, - byminute = 0, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) +non_trading_rules.append(good_friday) -# Recurrence rules for generating the market open/close for memorial -# day. Memorial day always occurs on the last monday of May. -memorial_day_opens = rrule.rrule( +memorial_day = rrule.rrule( rrule.MONTHLY, bymonth = 5, byweekday = (rrule.MO(-1)), - byhour = 14, - byminute = 30, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) -memorial_day_closes = rrule.rrule( - rrule.MONTHLY, - bymonth = 5, - byweekday = (rrule.MO(-1)), - byhour = 21, - byminute = 0, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) +non_trading_rules.append(memorial_day) -# Recurrence rules for generating the market open/close for July 4th. -july_4th_opens = rrule.rrule( +july_4th = rrule.rrule( rrule.MONTHLY, - bymonth = 6, + bymonth = 7, bymonthday = 4, - byhour = 14, - byminute = 30, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) -july_4th_closes = rrule.rrule( - rrule.MONTHLY, - bymonth = 6, - bymonthday = 4, - byhour = 21, - byminute = 0, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) +non_trading_rules.append(july_4th) -# Recurrence rule for generating the market open/close for labor day. -# Labor day is always the first monday of September. -labor_day_opens = rrule.rrule( +labor_day = rrule.rrule( rrule.MONTHLY, bymonth = 9, byweekday = (rrule.MO(1)), - byhour = 14, - byminute = 30, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) -labor_day_closes = rrule.rrule( - rrule.MONTHLY, - bymonth = 9, - byweekday = (rrule.MO(1)), - byhour = 21, - byminute = 0, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) +non_trading_rules.append(labor_day) -# Recurrence rule for generating the market open/close for -# thanksgiving. Thanksgiving always falls on the fourth thursday in -# November. (Who decides how these holidays work!?!) -thanksgiving_opens = rrule.rrule( +thanksgiving = rrule.rrule( rrule.MONTHLY, bymonth = 11, byweekday = (rrule.TH(-1)), - byhour = 14, - byminute = 30, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) -thanksgiving_closes = rrule.rrule( - rrule.MONTHLY, - bymonth = 11, - byweekday = (rrule.TH(-1)), - byhour = 21, - byminute = 0, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) +non_trading_rules.append(thanksgiving) -# Recurrence relation for generating the market open/close for -# christmas. Christmas always occurs on december 25th. - -christmas_opens = rrule.rrule( +christmas = rrule.rrule( rrule.MONTHLY, bymonth = 12, - bymonthday = 25, - byhour = 14, - byminute = 30, + bymonthday = 25, cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) -) -christmas_closes = rrule.rrule( - rrule.MONTHLY, - bymonth = 12, - bymonthday = 25, - byhour = 21, - byminute = 0, - cache = True, - dtstart = datetime(2000, 1,1,tzinfo = pytz.utc), - until=datetime(2014, 1, 1, tzinfo = pytz.utc) + dtstart = start, + until = end ) +non_trading_rules.append(christmas) -# All NYSE observed holidays. -holiday_opens = [ - new_years_opens, - mlk_opens, - presidents_day_opens, - good_friday_opens, - memorial_day_opens, - july_4th_opens, - labor_day_opens, - thanksgiving_opens, - christmas_opens -] -holiday_closes = [ - new_years_closes, - mlk_closes, - presidents_day_closes, - good_friday_closes, - memorial_day_closes, - july_4th_closes, - labor_day_closes, - thanksgiving_closes, - christmas_closes -] +non_trading_ruleset = rrule.rruleset() -# Valid market opens are given by all market opens minus holidays. -opens = rrule.rruleset(cache=True) -opens.rrule(market_opens_with_holidays) -for holiday_rule in holiday_opens: - opens.exrule(holiday_rule) - -closes = rrule.rruleset(cache=True) -closes.rrule(market_closes_with_holidays) -for holiday_rule in holiday_closes: - closes.exrule(holiday_rule) - -# This runs the calendar to load all data into a cache. -open_count = opens.count() -close_count = closes.count() +for rule in non_trading_rules: + non_trading_ruleset.rrule(rule) +non_trading_days = non_trading_ruleset.between(start, end, inc=True)