diff --git a/tests/utils/test_events.py b/tests/utils/test_events.py index c598ad69..90e2ebc0 100644 --- a/tests/utils/test_events.py +++ b/tests/utils/test_events.py @@ -344,76 +344,6 @@ 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, @@ -568,8 +498,7 @@ class TestStatelessRules(RuleTestCase): @subtest(param_range(MAX_MONTH_RANGE), 'n') def test_NDaysBeforeLastTradingDayOfMonth(self, n): should_trigger = partial( - NDaysBeforeLastTradingDayOfMonth(n, True).should_trigger, - env=self.env + NDaysBeforeLastTradingDayOfMonth(n).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 5952ea8a..4a9627fc 100644 --- a/zipline/utils/events.py +++ b/zipline/utils/events.py @@ -20,8 +20,6 @@ import datetime import pandas as pd import pytz -from zipline.errors import NoFurtherDataError - from .context_tricks import nop_context @@ -72,6 +70,29 @@ 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: @@ -413,28 +434,23 @@ class TradingDayOfWeekRule(six.with_metaclass(ABCMeta, StatelessRule)): raise NotImplementedError def calculate_start_and_end(self, dt, env): - while True: - new_date = self.date_func(dt, env) - if not new_date: - return + next_trading_day = _coerce_datetime( + env.add_trading_days( + self.td_delta, + self.date_func(dt, env), + ) + ) - try: - next_trading_day = env.add_trading_days( + # 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( self.td_delta, - new_date, + self.date_func(dt, env), ) - 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 @@ -445,11 +461,8 @@ 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. Return False if there is not trigger - # that occurs on or before the max day + # of the simulation start 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: @@ -474,16 +487,24 @@ class NthTradingDayOfWeek(TradingDayOfWeekRule): def get_first_trading_day_of_week(dt, env): prev = dt dt = env.previous_trading_day(dt) - # Traverse backward until we hit a week border, then jump back to the - # previous trading day. - while dt and dt.weekday() < prev.weekday(): + # 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(): prev = dt dt = env.previous_trading_day(dt) + if dt is None: + return prev if env.is_trading_day(prev): - return prev + return prev.date() else: - return env.next_trading_day(prev) + return env.next_trading_day(prev).date() date_func = get_first_trading_day_of_week @@ -501,68 +522,86 @@ 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 and dt.weekday() > prev.weekday(): + while dt.date().weekday() > prev.date().weekday(): prev = dt dt = env.next_trading_day(dt) if env.is_trading_day(prev): - return prev + return prev.date() else: - return env.previous_trading_day(prev) + return env.previous_trading_day(prev).date() date_func = get_last_trading_day_of_week -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): +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): + 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) - 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 + self.first_day = (dt if env.is_trading_day(dt) + else env.next_trading_day(dt)).date() + return self.first_day -class NDaysBeforeLastTradingDayOfMonth(TradingDayOfMonthRule): +class NDaysBeforeLastTradingDayOfMonth(StatelessRule): """ 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 @@ -575,12 +614,10 @@ class NDaysBeforeLastTradingDayOfMonth(TradingDayOfMonthRule): year = dt.year month = dt.month + 1 - last_day = env.previous_trading_day( + self.last_day = env.previous_trading_day( dt.replace(year=year, month=month, day=1) - ) - return last_day - - date_func = get_last_trading_day_of_month + ).date() + return self.last_day # Stateful rules @@ -634,11 +671,11 @@ class date_rules(object): @staticmethod def month_start(days_offset=0): - return NthTradingDayOfMonth(n=days_offset, invert=False) + return NthTradingDayOfMonth(n=days_offset) @staticmethod def month_end(days_offset=0): - return NDaysBeforeLastTradingDayOfMonth(n=days_offset, invert=True) + return NDaysBeforeLastTradingDayOfMonth(n=days_offset) @staticmethod def week_start(days_offset=0):