From 0cb4c387178001216b407540550b46294380374a Mon Sep 17 00:00:00 2001 From: llllllllll Date: Tue, 10 Nov 2015 15:09:35 -0500 Subject: [PATCH] ENH: Allow users to pass a context manager to wrap all scheduled functions. This includes handle_data. --- tests/test_algorithm.py | 60 +++++++++++++++++++++++++++++++-- zipline/algorithm.py | 4 ++- zipline/utils/context_tricks.py | 60 +++++++++++++++++++++++++++++++++ zipline/utils/events.py | 37 +++++++++++++++++--- 4 files changed, 153 insertions(+), 8 deletions(-) create mode 100644 zipline/utils/context_tricks.py diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index ed3a298d..074d947f 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -79,7 +79,7 @@ from zipline.test_algorithms import ( record_float_magic, record_variables, ) - +from zipline.utils.context_tricks import CallbackManager import zipline.utils.events from zipline.utils.test_utils import ( assert_single_position, @@ -96,7 +96,7 @@ from zipline.assets import Equity from zipline.finance.execution import LimitOrder from zipline.finance.trading import SimulationParameters from zipline.utils.api_support import set_algo_instance -from zipline.utils.events import DateRuleFactory, TimeRuleFactory +from zipline.utils.events import DateRuleFactory, TimeRuleFactory, Always from zipline.algorithm import TradingAlgorithm from zipline.protocol import DATASOURCE_TYPE from zipline.finance.trading import TradingEnvironment @@ -347,6 +347,62 @@ class TestMiscellaneousAPI(TestCase): self.assertEqual(algo.func_called, algo.days) + def test_event_context(self): + expected_data = [] + collected_data_pre = [] + collected_data_post = [] + function_stack = [] + + def pre(data): + function_stack.append(pre) + collected_data_pre.append(data) + + def post(data): + function_stack.append(post) + collected_data_post.append(data) + + def initialize(context): + context.add_event(Always(), f) + context.add_event(Always(), g) + + def handle_data(context, data): + function_stack.append(handle_data) + expected_data.append(data) + + def f(context, data): + function_stack.append(f) + + def g(context, data): + function_stack.append(g) + + algo = TradingAlgorithm( + initialize=initialize, + handle_data=handle_data, + sim_params=self.sim_params, + create_event_context=CallbackManager(pre, post), + env=self.env, + ) + algo.run(self.source) + + self.assertEqual(len(expected_data), 779) + self.assertEqual(collected_data_pre, expected_data) + self.assertEqual(collected_data_post, expected_data) + + self.assertEqual( + len(function_stack), + 779 * 5, + 'Incorrect number of functions called: %s != 779' % + len(function_stack), + ) + expected_functions = [pre, handle_data, f, g, post] * 779 + for n, (f, g) in enumerate(zip(function_stack, expected_functions)): + self.assertEqual( + f, + g, + 'function at position %d was incorrect, expected %s but got %s' + % (n, g.__name__, f.__name__), + ) + @parameterized.expand([ ('daily',), ('minute'), diff --git a/zipline/algorithm.py b/zipline/algorithm.py index c2e3a628..b6507e30 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -270,7 +270,9 @@ class TradingAlgorithm(object): self._before_trading_start = None self._analyze = None - self.event_manager = EventManager() + self.event_manager = EventManager( + create_context=kwargs.pop('create_event_context', None), + ) if self.algoscript is not None: filename = kwargs.pop('algo_filename', None) diff --git a/zipline/utils/context_tricks.py b/zipline/utils/context_tricks.py new file mode 100644 index 00000000..6d76879c --- /dev/null +++ b/zipline/utils/context_tricks.py @@ -0,0 +1,60 @@ +from contextlib import contextmanager + + +def _nop(*args, **kwargs): + pass + + +class CallbackManager(object): + """Create a context manager from a pre-execution callback and a + post-execution callback. + + Parameters + ---------- + pre : (...) -> any, optional + A pre-execution callback. This will be passed ``*args`` and + ``**kwargs``. + post : (...) -> any, optional + A post-execution callback. This will be passed ``*args`` and + ``**kwargs``. + + Notes + ----- + The enter value of this context manager will be the result of calling + ``pre(*args, **kwargs)`` + + Examples + -------- + >>> def pre(where): + ... print('entering %s block' % where) + >>> def post(where): + ... print('exiting %s block' % where) + >>> manager = CallbackManager(pre, post) + >>> with manager('example'): + ... print('inside example block') + entering example block + inside example + exiting example block + + These are reusable with different args: + >>> with manager('another'): + ... print('inside another block') + entering another block + inside another block + exiting another block + """ + def __init__(self, pre=None, post=None): + pre = pre if pre is not None else _nop + post = post if post is not None else _nop + + @contextmanager + def _callback_manager_context(*args, **kwargs): + try: + yield pre(*args, **kwargs) + finally: + post(*args, **kwargs) + + self._callback_manager_context = _callback_manager_context + + def __call__(self, *args, **kwargs): + return self._callback_manager_context(*args, **kwargs) diff --git a/zipline/utils/events.py b/zipline/utils/events.py index 55197698..a9cd6744 100644 --- a/zipline/utils/events.py +++ b/zipline/utils/events.py @@ -169,14 +169,35 @@ def _build_time(time, kwargs): return datetime.time(**kwargs) -class EventManager(object): +@object.__new__ +class _nop_context(object): + """A nop context manager. """ - Manages a list of Event objects. + def __enter__(self): + pass + + def __exit__(self, *excinfo): + pass + + +class EventManager(object): + """Manages a list of Event objects. This manages the logic for checking the rules and dispatching to the handle_data function of the Events. + + Parameters + ---------- + create_context : (BarData) -> context manager, optional + An optional callback to produce a context manager to wrap the calls + to handle_data. This will be passed the current BarData. """ - def __init__(self): + def __init__(self, create_context=None): self._events = [] + self._create_context = ( + create_context + if create_context is not None else + lambda *_: _nop_context + ) def add_event(self, event, prepend=False): """ @@ -188,8 +209,14 @@ class EventManager(object): self._events.append(event) def handle_data(self, context, data, dt): - for event in self._events: - event.handle_data(context, data, dt, context.trading_environment) + with self._create_context(data): + for event in self._events: + event.handle_data( + context, + data, + dt, + context.trading_environment, + ) class Event(namedtuple('Event', ['rule', 'callback'])):