Merge pull request #1226 from quantopian/schedule_function_resilience

MAINT: Make schedule function rules resilient for no trading day

tests/utils/test_events.py

@@ -344,6 +344,76 @@ class TestStatelessRules(RuleTestCase):
             )
         )
 
+    @parameterized.expand([
+        # No more trading days starting from the new week
+        ('2016-01-08', '2016-01-11', '2016-01-04', None, 0, True),
+        # New week has a trading day, but no more after that day
+        ('2016-01-11', '2016-01-11', '2016-01-05', '2016-01-11', 1, True),
+        ('2016-01-11', '2016-01-11', '2016-01-07', '2016-01-11', 3, True),
+        # Has not had a previous trigger calculated yet
+        ('2016-01-11', '2016-01-11', None, '2016-01-11', 3, False),
+    ])
+    def test_NthTradingDayOfWeek_no_days_left(self, max_date, dt,
+                                              expected_trigger,
+                                              first_trading_day_of_week,
+                                              offset,
+                                              has_previous_trigger):
+        """
+        Test that when we try to calculate the first trading day of week but
+        there are no trading days left going forward,
+        we don't crash and instead don't update the trigger
+        """
+        max_date = pd.Timestamp(max_date, tz='UTC')
+        dt = pd.Timestamp(dt, tz='UTC')
+        env = TradingEnvironment(max_date=max_date)
+        rule = NthTradingDayOfWeek(offset)
+
+        if has_previous_trigger:
+            next_midnight_timestamp = pd.Timestamp(expected_trigger, tz='UTC')
+            next_date_start, next_date_end = \
+                env.get_open_and_close(next_midnight_timestamp)
+        else:
+            next_midnight_timestamp = next_date_start = next_date_end = None
+
+        rule.next_midnight_timestamp = next_midnight_timestamp
+        rule.next_date_start = next_date_start
+        rule.next_date_end = next_date_end
+
+        expected_first_trading_day_of_week = \
+            pd.Timestamp(first_trading_day_of_week, tz='UTC') \
+            if first_trading_day_of_week else None
+
+        self.assertEqual(rule.get_first_trading_day_of_week(dt, env),
+                         expected_first_trading_day_of_week)
+        self.assertEqual(rule.should_trigger(dt, env), False)
+        self.assertEqual(rule.next_date_end, next_date_end)
+        self.assertEqual(rule.next_date_start, next_date_start)
+        self.assertEqual(rule.next_midnight_timestamp, next_midnight_timestamp)
+
+    @parameterized.expand([
+        ('2016-01-04', '2016-01-04', 1), ('2016-01-04', '2016-01-04', 2),
+    ])
+    def test_NthTradingDayOfMonth_no_days_left(self, max_date, dt, offset):
+        max_date = pd.Timestamp(max_date, tz='UTC')
+        dt = pd.Timestamp(dt, tz='UTC')
+        env = TradingEnvironment(max_date=max_date)
+        rule = NthTradingDayOfMonth(offset)
+
+        self.assertEqual(rule.get_trigger_day_of_month(dt, env), None)
+
+    @parameterized.expand([
+        ('2016-01-29', '2016-01-29', 1), ('2016-01-29', '2016-01-29', 2),
+    ])
+    def test_NDaysBeforeLastTradingDayOfMonth_no_days_left(self, min_date, dt,
+                                                           offset):
+        min_date = pd.Timestamp(min_date, tz='UTC')
+        dt = pd.Timestamp(dt, tz='UTC')
+        env = TradingEnvironment(min_date=min_date)
+        rule = NDaysBeforeLastTradingDayOfMonth(offset, True)
+
+        self.assertEqual(rule.get_trigger_day_of_month(dt, env),
+                         None)
+
     @subtest(param_range(MAX_WEEK_RANGE), 'n')
     def test_NthTradingDayOfWeek(self, n):
         should_trigger = partial(NthTradingDayOfWeek(n).should_trigger,
@@ -498,7 +568,8 @@ class TestStatelessRules(RuleTestCase):
     @subtest(param_range(MAX_MONTH_RANGE), 'n')
     def test_NDaysBeforeLastTradingDayOfMonth(self, n):
         should_trigger = partial(
-            NDaysBeforeLastTradingDayOfMonth(n).should_trigger, env=self.env
+            NDaysBeforeLastTradingDayOfMonth(n, True).should_trigger,
+            env=self.env
         )
         for n_days_before, d in enumerate(reversed(self.sept_days)):
             for m in self.env.market_minutes_for_day(d):

zipline/utils/events.py

@@ -20,6 +20,8 @@ import datetime
 import pandas as pd
 import pytz
 
+from zipline.errors import NoFurtherDataError
+
 from .context_tricks import nop_context
 
 
@@ -70,29 +72,6 @@ def ensure_utc(time, tz='UTC'):
     return time.replace(tzinfo=pytz.utc)
 
 
-def _coerce_datetime(maybe_dt):
-    if isinstance(maybe_dt, datetime.datetime):
-        return maybe_dt
-    elif isinstance(maybe_dt, datetime.date):
-        return datetime.datetime(
-            year=maybe_dt.year,
-            month=maybe_dt.month,
-            day=maybe_dt.day,
-            tzinfo=pytz.utc,
-        )
-    elif isinstance(maybe_dt, (tuple, list)) and len(maybe_dt) == 3:
-        year, month, day = maybe_dt
-        return datetime.datetime(
-            year=year,
-            month=month,
-            day=day,
-            tzinfo=pytz.utc,
-        )
-    else:
-        raise TypeError('Cannot coerce %s into a datetime.datetime'
-                        % type(maybe_dt).__name__)
-
-
 def _out_of_range_error(a, b=None, var='offset'):
     start = 0
     if b is None:
@@ -434,23 +413,28 @@ class TradingDayOfWeekRule(six.with_metaclass(ABCMeta, StatelessRule)):
         raise NotImplementedError
 
     def calculate_start_and_end(self, dt, env):
-        next_trading_day = _coerce_datetime(
-            env.add_trading_days(
-                self.td_delta,
-                self.date_func(dt, env),
-            )
-        )
+        while True:
+            new_date = self.date_func(dt, env)
+            if not new_date:
+                return
 
-        # 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.isocalendar()[1] != dt.isocalendar()[1]:
-            dt += datetime.timedelta(days=7)
-            next_trading_day = _coerce_datetime(
-                env.add_trading_days(
+            try:
+                next_trading_day = env.add_trading_days(
                     self.td_delta,
-                    self.date_func(dt, env),
+                    new_date,
                 )
-            )
+            except NoFurtherDataError:
+                return
+
+            if not next_trading_day:
+                return
+
+            # 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
+            if next_trading_day.isocalendar()[1] == dt.isocalendar()[1]:
+                break
+            else:
+                dt += datetime.timedelta(days=7)
 
         next_open, next_close = env.get_open_and_close(next_trading_day)
         self.next_date_start = next_open
@@ -461,8 +445,11 @@ class TradingDayOfWeekRule(six.with_metaclass(ABCMeta, StatelessRule)):
         if self.next_date_start is None:
             # First time this method has been called. Calculate the midnight,
             # open, and close for the first trigger, which occurs on the week
-            # of the simulation start
+            # of the simulation start. Return False if there is not trigger
+            # that occurs on or before the max day
             self.calculate_start_and_end(dt, env)
+            if not self.next_date_start:
+                return False
 
         # If we've passed the trigger, calculate the next one
         if dt > self.next_date_end:
@@ -487,24 +474,16 @@ class NthTradingDayOfWeek(TradingDayOfWeekRule):
     def get_first_trading_day_of_week(dt, env):
         prev = dt
         dt = env.previous_trading_day(dt)
-        # If we're on the first trading day of the TradingEnvironment,
-        # calling previous_trading_day on it will return None, which
-        # will blow up when we try and call .date() on it. The first
-        # trading day of the env is also the first trading day of the
-        # week(in the TradingEnvironment, at least), so just return
-        # that date.
-        if dt is None:
-            return prev
-        while dt.date().weekday() < prev.date().weekday():
+        # Traverse backward until we hit a week border, then jump back to the
+        # previous trading day.
+        while dt and dt.weekday() < prev.weekday():
             prev = dt
             dt = env.previous_trading_day(dt)
-            if dt is None:
-                return prev
 
         if env.is_trading_day(prev):
-            return prev.date()
+            return prev
         else:
-            return env.next_trading_day(prev).date()
+            return env.next_trading_day(prev)
 
     date_func = get_first_trading_day_of_week
 
@@ -522,86 +501,68 @@ class NDaysBeforeLastTradingDayOfWeek(TradingDayOfWeekRule):
         dt = env.next_trading_day(dt)
         # Traverse forward until we hit a week border, then jump back to the
         # previous trading day.
-        while dt.date().weekday() > prev.date().weekday():
+        while dt and dt.weekday() > prev.weekday():
             prev = dt
             dt = env.next_trading_day(dt)
 
         if env.is_trading_day(prev):
-            return prev.date()
+            return prev
         else:
-            return env.previous_trading_day(prev).date()
+            return env.previous_trading_day(prev)
 
     date_func = get_last_trading_day_of_week
 
 
-class NthTradingDayOfMonth(StatelessRule):
+class TradingDayOfMonthRule(six.with_metaclass(ABCMeta, StatelessRule)):
+    def __init__(self, n=0, invert=False):
+        if not 0 <= n < MAX_MONTH_RANGE:
+            raise _out_of_range_error(MAX_MONTH_RANGE)
+        self.month = None
+        self.date = None
+        self.td_delta = -n if invert else n
+
+    def should_trigger(self, dt, env):
+        return self.get_trigger_day_of_month(dt, env) == env.normalize_date(dt)
+
+    @abstractmethod
+    def date_func(self, dt, env):
+        raise NotImplementedError
+
+    def get_trigger_day_of_month(self, dt, env):
+        if self.month == dt.month:
+            # We already computed the day for this month.
+            return self.date
+
+        self.date = self.date_func(dt, env)
+        if self.td_delta and self.date:
+            try:
+                self.date = env.add_trading_days(self.td_delta, self.date)
+            except NoFurtherDataError:
+                self.date = None
+
+        return self.date
+
+
+class NthTradingDayOfMonth(TradingDayOfMonthRule):
     """
     A rule that triggers on the nth trading day of the month.
     This is zero-indexed, n=0 is the first trading day of the month.
     """
-    def __init__(self, n=0):
-        if not 0 <= n < MAX_MONTH_RANGE:
-            raise _out_of_range_error(MAX_MONTH_RANGE)
-        self.td_delta = n
-        self.month = None
-        self.day = None
-
-    def should_trigger(self, dt, env):
-        return self.get_nth_trading_day_of_month(dt, env) == dt.date()
-
-    def get_nth_trading_day_of_month(self, dt, env):
-        if self.month == dt.month:
-            # We already computed the day for this month.
-            return self.day
-
-        if not self.td_delta:
-            self.day = self.get_first_trading_day_of_month(dt, env)
-        else:
-            self.day = env.add_trading_days(
-                self.td_delta,
-                self.get_first_trading_day_of_month(dt, env),
-            ).date()
-
-        return self.day
-
     def get_first_trading_day_of_month(self, dt, env):
         self.month = dt.month
 
         dt = dt.replace(day=1)
-        self.first_day = (dt if env.is_trading_day(dt)
-                          else env.next_trading_day(dt)).date()
-        return self.first_day
+        first_day = (dt if env.is_trading_day(dt)
+                     else env.next_trading_day(dt))
+        return first_day
+
+    date_func = get_first_trading_day_of_month
 
 
-class NDaysBeforeLastTradingDayOfMonth(StatelessRule):
+class NDaysBeforeLastTradingDayOfMonth(TradingDayOfMonthRule):
     """
     A rule that triggers n days before the last trading day of the month.
     """
-    def __init__(self, n=0):
-        if not 0 <= n < MAX_MONTH_RANGE:
-            raise _out_of_range_error(MAX_MONTH_RANGE)
-        self.td_delta = -n
-        self.month = None
-        self.day = None
-
-    def should_trigger(self, dt, env):
-        return self.get_nth_to_last_trading_day_of_month(dt, env) == dt.date()
-
-    def get_nth_to_last_trading_day_of_month(self, dt, env):
-        if self.month == dt.month:
-            # We already computed the last day for this month.
-            return self.day
-
-        if not self.td_delta:
-            self.day = self.get_last_trading_day_of_month(dt, env)
-        else:
-            self.day = env.add_trading_days(
-                self.td_delta,
-                self.get_last_trading_day_of_month(dt, env),
-            ).date()
-
-        return self.day
-
     def get_last_trading_day_of_month(self, dt, env):
         self.month = dt.month
 
@@ -614,10 +575,12 @@ class NDaysBeforeLastTradingDayOfMonth(StatelessRule):
             year = dt.year
             month = dt.month + 1
 
-        self.last_day = env.previous_trading_day(
+        last_day = env.previous_trading_day(
             dt.replace(year=year, month=month, day=1)
-        ).date()
-        return self.last_day
+        )
+        return last_day
+
+    date_func = get_last_trading_day_of_month
 
 
 # Stateful rules
@@ -671,11 +634,11 @@ class date_rules(object):
 
     @staticmethod
     def month_start(days_offset=0):
-        return NthTradingDayOfMonth(n=days_offset)
+        return NthTradingDayOfMonth(n=days_offset, invert=False)
 
     @staticmethod
     def month_end(days_offset=0):
-        return NDaysBeforeLastTradingDayOfMonth(n=days_offset)
+        return NDaysBeforeLastTradingDayOfMonth(n=days_offset, invert=True)
 
     @staticmethod
     def week_start(days_offset=0):