BUG: week_end rule with offset=0 skips every other week

tests/utils/test_events.py

@@ -18,6 +18,7 @@ from functools import partial
 from inspect import isabstract
 import random
 from unittest import TestCase
+from datetime import timedelta
 
 from nose_parameterized import parameterized
 import pandas as pd
@@ -25,7 +26,7 @@ from six import iteritems
 from six.moves import range, map
 
 from zipline.finance.trading import TradingEnvironment
-from zipline.testing import subtest
+from zipline.testing import subtest, parameter_space
 import zipline.utils.events
 from zipline.utils.events import (
     EventRule,
@@ -347,34 +348,79 @@ class TestStatelessRules(RuleTestCase):
     def test_NthTradingDayOfWeek(self, n):
         should_trigger = partial(NthTradingDayOfWeek(n).should_trigger,
                                  env=self.env)
-        prev_day = self.sept_week[0].date()
-        n_tdays = 0
+        first_of_week = self.sept_week[0]
+        n_triggered = 0
         for m in self.sept_week:
-            if prev_day < m.date():
-                n_tdays += 1
-            prev_day = m.date()
-
             if should_trigger(m):
-                self.assertEqual(n_tdays, n)
-            else:
-                self.assertNotEqual(n_tdays, n)
+                n_triggered += 1
+                self.assertEqual(m, first_of_week + timedelta(days=n))
+        self.assertEqual(n_triggered, 1)
 
     @subtest(param_range(MAX_WEEK_RANGE), 'n')
     def test_NDaysBeforeLastTradingDayOfWeek(self, n):
         should_trigger = partial(
             NDaysBeforeLastTradingDayOfWeek(n).should_trigger, env=self.env
         )
+        first_of_week = self.sept_week[0]
+        n_triggered = 0
         for m in self.sept_week:
             if should_trigger(m):
-                n_tdays = 0
-                date = m.to_datetime().date()
-                next_date = self.env.next_trading_day(date)
-                while next_date.weekday() > date.weekday():
-                    date = next_date
-                    next_date = self.env.next_trading_day(date)
-                    n_tdays += 1
+                n_triggered += 1
+                self.assertEqual(m, first_of_week + timedelta(days=4-n))
+        self.assertEqual(n_triggered, 1)
 
-                self.assertEqual(n_tdays, n)
+    @parameter_space(n=(0, 1, 2, 3, 4), type=('week_start', 'week_end'))
+    def test_edge_cases_for_TradingDayOfWeek(self, n, type):
+        """
+        Test that we account for midweek holidays. Monday 01/20 is a holiday.
+        Ensure that the trigger date for that week is adjustmented
+        appropriately, or thrown out if not enough trading days.
+        """
+        jan_minutes = self.env.minutes_for_days_in_range(
+            datetime.date(year=2014, month=1, day=6),
+            datetime.date(year=2014, month=1, day=31)
+        )
+
+        if type == 'week_start':
+            rule = NthTradingDayOfWeek
+            # Expect to trigger on the first trading day of the week, plus the
+            # offset
+            trigger_dates = [
+                pd.Timestamp('2014-01-06 14:31:00', tz='UTC'),
+                pd.Timestamp('2014-01-13 14:31:00', tz='UTC'),
+                pd.Timestamp('2014-01-21 14:31:00', tz='UTC'),
+                pd.Timestamp('2014-01-27 14:31:00', tz='UTC'),
+            ]
+            trigger_dates = [x + timedelta(days=n) for x in trigger_dates]
+        else:
+            rule = NDaysBeforeLastTradingDayOfWeek
+            # Expect to trigger on the last trading day of the week, minus the
+            # offset
+            trigger_dates = [
+                pd.Timestamp('2014-01-10 14:31:00', tz='UTC'),
+                pd.Timestamp('2014-01-17 14:31:00', tz='UTC'),
+                pd.Timestamp('2014-01-24 14:31:00', tz='UTC'),
+                pd.Timestamp('2014-01-31 14:31:00', tz='UTC'),
+            ]
+            trigger_dates = [x - timedelta(days=n) for x in trigger_dates]
+
+        should_trigger = partial(
+            rule(n).should_trigger, env=self.env
+        )
+
+        # If offset is 4, there is not enough trading days in the short week,
+        # and so it should not trigger
+        if n == 4:
+            del trigger_dates[2]
+        trigger_dates = iter(trigger_dates)
+
+        n_triggered = 0
+        for m in jan_minutes:
+            if should_trigger(m):
+                self.assertEqual(m, next(trigger_dates))
+                n_triggered += 1
+
+        self.assertEqual(n_triggered, 4 if n != 4 else 3)
 
     @subtest(param_range(MAX_MONTH_RANGE), 'n')
     def test_NthTradingDayOfMonth(self, n):

zipline/utils/events.py

@@ -443,6 +443,18 @@ class TradingDayOfWeekRule(six.with_metaclass(ABCMeta, StatelessRule)):
             )
         )
 
+        # If after applying the offset to the start/end day of the week, we get
+        # day in a different week, skip this week and go on to the next
+        while next_trading_day.date().isocalendar()[1] != \
+                dt.date().isocalendar()[1]:
+            dt += datetime.timedelta(days=7)
+            next_trading_day = _coerce_datetime(
+                env.add_trading_days(
+                    self.td_delta,
+                    self.date_func(dt, env),
+                )
+            )
+
         next_open, next_close = env.get_open_and_close(next_trading_day)
         self.next_date_start = next_open
         self.next_date_end = next_close
@@ -455,14 +467,12 @@ class TradingDayOfWeekRule(six.with_metaclass(ABCMeta, StatelessRule)):
             # open, and close of the next matching day.
             self.calculate_start_and_end(dt, env)
 
-        # if the next matching day is in the past, calculate the next one.
-        if dt > self.next_date_end:
-            self.calculate_start_and_end(dt + datetime.timedelta(days=7),
-                                         env)
-
-        # if the given dt is within the next matching day, return true.
+        # if the given dt is within the next matching day, return true. Also
+        # calculate the start and end dates for the next trigger
         if self.next_date_start <= dt <= self.next_date_end or \
                 dt == self.next_midnight_timestamp:
+            self.calculate_start_and_end(dt + datetime.timedelta(days=7),
+                                         env)
             return True
 
         return False
@@ -490,7 +500,11 @@ class NthTradingDayOfWeek(TradingDayOfWeekRule):
             dt = env.previous_trading_day(dt)
             if dt is None:
                 return prev
-        return prev.date()
+
+        if env.is_trading_day(prev):
+            return prev.date()
+        else:
+            return env.next_trading_day(prev).date()
 
     date_func = get_first_trading_day_of_week
 
@@ -511,7 +525,11 @@ class NDaysBeforeLastTradingDayOfWeek(TradingDayOfWeekRule):
         while dt.date().weekday() > prev.date().weekday():
             prev = dt
             dt = env.next_trading_day(dt)
-        return prev.date()
+
+        if env.is_trading_day(prev):
+            return prev.date()
+        else:
+            return env.previous_trading_day(prev).date()
 
     date_func = get_last_trading_day_of_week