From 93fa848c19340618f3dd4d71a06b100025b033b9 Mon Sep 17 00:00:00 2001 From: Richard Frank Date: Tue, 26 Dec 2017 11:12:25 -0500 Subject: [PATCH] BUG: cal attribute of aggregate rules was not being propagated --- tests/events/test_events.py | 8 +++--- tests/test_algorithm.py | 36 ++++++++++++++++++-------- zipline/algorithm.py | 2 +- zipline/utils/events.py | 50 ++++++++++++++++++++++++++----------- 4 files changed, 65 insertions(+), 31 deletions(-) diff --git a/tests/events/test_events.py b/tests/events/test_events.py index d5573c60d1..bebc044555 100644 --- a/tests/events/test_events.py +++ b/tests/events/test_events.py @@ -136,8 +136,8 @@ def test_build_time_kwargs(self): class TestEventManager(TestCase): def setUp(self): self.em = EventManager() - self.event1 = Event(Always(), lambda context, data: None) - self.event2 = Event(Always(), lambda context, data: None) + self.event1 = Event(Always()) + self.event2 = Event(Always()) def test_add_event(self): self.em.add_event(self.event1) @@ -162,9 +162,7 @@ def should_trigger(self, dt): return True for r in [CountingRule] * 5: - self.em.add_event( - Event(r(), lambda context, data: None) - ) + self.em.add_event(Event(r())) self.em.handle_data(None, None, datetime.datetime.now()) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index c8a783938e..967e6d8566 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -34,6 +34,7 @@ import pytz from pandas.core.common import PerformanceWarning +import zipline.api from zipline import run_algorithm from zipline import TradingAlgorithm from zipline.api import FixedSlippage @@ -181,8 +182,14 @@ from zipline.utils.api_support import ZiplineAPI, set_algo_instance from zipline.utils.calendars import get_calendar, register_calendar from zipline.utils.context_tricks import CallbackManager, nop_context -import zipline.utils.events -from zipline.utils.events import date_rules, time_rules, Always +from zipline.utils.events import ( + date_rules, + time_rules, + Always, + ComposedRule, + Never, + OncePerDay, +) import zipline.utils.factory as factory # Because test cases appear to reuse some resources. @@ -635,27 +642,36 @@ def nop(*args, **kwargs): ) # Schedule something for NOT Always. - algo.schedule_function(nop, time_rule=zipline.utils.events.Never()) + # Compose two rules to ensure calendar is set properly. + algo.schedule_function(nop, time_rule=Never() & Always()) event_rule = algo.event_manager._events[1].rule - - self.assertIsInstance(event_rule, zipline.utils.events.OncePerDay) + self.assertIsInstance(event_rule, OncePerDay) + self.assertEqual(event_rule.cal, algo.trading_calendar) inner_rule = event_rule.rule - self.assertIsInstance(inner_rule, zipline.utils.events.ComposedRule) + self.assertIsInstance(inner_rule, ComposedRule) + self.assertEqual(inner_rule.cal, algo.trading_calendar) first = inner_rule.first second = inner_rule.second composer = inner_rule.composer - self.assertIsInstance(first, zipline.utils.events.Always) + self.assertIsInstance(first, Always) + self.assertEqual(first.cal, algo.trading_calendar) + self.assertEqual(second.cal, algo.trading_calendar) if mode == 'daily': - self.assertIsInstance(second, zipline.utils.events.Always) + self.assertIsInstance(second, Always) else: - self.assertIsInstance(second, zipline.utils.events.Never) + self.assertIsInstance(second, ComposedRule) + self.assertIsInstance(second.first, Never) + self.assertEqual(second.first.cal, algo.trading_calendar) + + self.assertIsInstance(second.second, Always) + self.assertEqual(second.second.cal, algo.trading_calendar) - self.assertIs(composer, zipline.utils.events.ComposedRule.lazy_and) + self.assertIs(composer, ComposedRule.lazy_and) def test_asset_lookup(self): algo = TradingAlgorithm(env=self.env) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 07dcbd3508..d40e9ac729 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -1064,7 +1064,7 @@ def fetch_csv(self, return csv_data_source - def add_event(self, rule=None, callback=None): + def add_event(self, rule, callback): """Adds an event to the algorithm's EventManager. Parameters diff --git a/zipline/utils/events.py b/zipline/utils/events.py index 878f1b0088..de161e9c15 100644 --- a/zipline/utils/events.py +++ b/zipline/utils/events.py @@ -223,7 +223,7 @@ class Event(namedtuple('Event', ['rule', 'callback'])): with the current algorithm context, data, and datetime only when the rule is triggered. """ - def __new__(cls, rule=None, callback=None): + def __new__(cls, rule, callback=None): callback = callback or (lambda *args, **kwargs: None) return super(cls, cls).__new__(cls, rule=rule, callback=callback) @@ -236,6 +236,18 @@ def handle_data(self, context, data, dt): class EventRule(six.with_metaclass(ABCMeta)): + # Instances of EventRule are assigned a calendar instance when scheduling + # a function. + _cal = None + + @property + def cal(self): + return self._cal + + @cal.setter + def cal(self, value): + self._cal = value + @abstractmethod def should_trigger(self, dt): """ @@ -301,6 +313,15 @@ def lazy_and(first_should_trigger, second_should_trigger, dt): """ return first_should_trigger(dt) and second_should_trigger(dt) + @property + def cal(self): + return self.first.cal + + @cal.setter + def cal(self, value): + # Thread the calendar through to the underlying rules. + self.first.cal = self.second.cal = value + class Always(StatelessRule): """ @@ -546,11 +567,14 @@ class StatefulRule(EventRule): def __init__(self, rule=None): self.rule = rule or Always() - def new_should_trigger(self, callable_): - """ - Replace the should trigger implementation for the current rule. - """ - self.should_trigger = callable_ + @property + def cal(self): + return self.rule.cal + + @cal.setter + def cal(self, value): + # Thread the calendar through to the underlying rule. + self.rule.cal = value class OncePerDay(StatefulRule): @@ -614,16 +638,12 @@ def make_eventrule(date_rule, time_rule, cal, half_days=True): """ Constructs an event rule from the factory api. """ - - # Insert the calendar in to the individual rules - date_rule.cal = cal - time_rule.cal = cal - if half_days: inner_rule = date_rule & time_rule else: - nhd_rule = NotHalfDay() - nhd_rule.cal = cal - inner_rule = date_rule & time_rule & nhd_rule + inner_rule = date_rule & time_rule & NotHalfDay() - return OncePerDay(rule=inner_rule) + opd = OncePerDay(rule=inner_rule) + # This is where a scheduled function's rule is associated with a calendar. + opd.cal = cal + return opd