From 7b82b9a2fde79e10827c6690a4640d1b626fda6b Mon Sep 17 00:00:00 2001 From: Andrew Liang Date: Mon, 9 May 2016 11:52:59 -0400 Subject: [PATCH] MAINT: Make schedule function rules resilient for no trading day Make the rules resilient to env.previous_trading_day or env.next_trading_day being None --- tests/utils/test_events.py | 73 ++++++++++++++- zipline/utils/events.py | 177 +++++++++++++++---------------------- 2 files changed, 142 insertions(+), 108 deletions(-) diff --git a/tests/utils/test_events.py b/tests/utils/test_events.py index 90e2ebc056..c598ad69f6 100644 --- a/tests/utils/test_events.py +++ b/tests/utils/test_events.py @@ -344,6 +344,76 @@ def test_NthTradingDayOfWeek_day_zero(self): ) ) + @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 @@ def test_NthTradingDayOfMonth(self, n): @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): diff --git a/zipline/utils/events.py b/zipline/utils/events.py index 4a9627fce9..5952ea8a97 100644 --- a/zipline/utils/events.py +++ b/zipline/utils/events.py @@ -20,6 +20,8 @@ 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 @@ def date_func(self, dt, env): 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 @@ def should_trigger(self, dt, env): 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 @@ def get_last_trading_day_of_week(dt, env): 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): - """ - 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): +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.td_delta = n self.month = None - self.day = None + self.date = None + self.td_delta = -n if invert else n def should_trigger(self, dt, env): - return self.get_nth_trading_day_of_month(dt, env) == dt.date() + return self.get_trigger_day_of_month(dt, env) == env.normalize_date(dt) - def get_nth_trading_day_of_month(self, dt, env): + @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.day + return self.date - 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() + 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 - return self.day +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 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 @@ def get_last_trading_day_of_month(self, dt, env): 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):