From 93429d69f50bd576fc6185bce3028c46df4f688d Mon Sep 17 00:00:00 2001 From: Joe Jevnik Date: Wed, 5 Nov 2014 13:02:10 -0500 Subject: [PATCH] BUG: HistoryContainer creation at runtime did not work as intended. --- tests/test_algorithm.py | 54 +++++++++++++++++++++++----- tests/test_history.py | 6 ++-- tests/utils/test_events.py | 3 +- zipline/algorithm.py | 16 ++++++--- zipline/history/history_container.py | 34 ++++++++++-------- 5 files changed, 82 insertions(+), 31 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 9c1075d9..f7df258c 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -714,6 +714,18 @@ def handle_data(context, data): class TestHistory(TestCase): + @classmethod + def setUpClass(cls): + cls._start = pd.Timestamp('1991-01-01', tz='UTC') + cls._end = pd.Timestamp('1991-01-15', tz='UTC') + cls.sim_params = factory.create_simulation_parameters( + data_frequency='minute', + ) + + @property + def source(self): + return RandomWalkSource(start=self._start, end=self._end) + def test_history(self): history_algo = """ from zipline.api import history, add_history @@ -724,16 +736,42 @@ def initialize(context): def handle_data(context, data): df = history(10, '1d', 'price') """ - start = pd.Timestamp('1991-01-01', tz='UTC') - end = pd.Timestamp('1991-01-15', tz='UTC') - source = RandomWalkSource(start=start, - end=end) - sim_params = factory.create_simulation_parameters( - data_frequency='minute') - algo = TradingAlgorithm(script=history_algo, sim_params=sim_params) - output = algo.run(source) + + algo = TradingAlgorithm( + script=history_algo, + sim_params=self.sim_params, + ) + output = algo.run(self.source) self.assertIsNot(output, None) + def test_history_without_add(self): + def handle_data(algo, data): + algo.history(1, '1m', 'price') + + algo = TradingAlgorithm( + initialize=lambda _: None, + handle_data=handle_data, + sim_params=self.sim_params, + ) + algo.run(self.source) + + self.assertIsNotNone(algo.history_container) + self.assertEqual(algo.history_container.buffer_panel.window_length, 1) + + def test_add_history_in_handle_data(self): + def handle_data(algo, data): + algo.add_history(1, '1m', 'price') + + algo = TradingAlgorithm( + initialize=lambda _: None, + handle_data=handle_data, + sim_params=self.sim_params, + ) + algo.run(self.source) + + self.assertIsNotNone(algo.history_container) + self.assertEqual(algo.history_container.buffer_panel.window_length, 1) + class TestGetDatetime(TestCase): diff --git a/tests/test_history.py b/tests/test_history.py index c95c9837..8db6d495 100644 --- a/tests/test_history.py +++ b/tests/test_history.py @@ -997,7 +997,7 @@ class TestHistoryContainerResize(TestCase): ) for spec in to_add: - container.ensure_spec(spec, initial_dt) + container.ensure_spec(spec, initial_dt, bar_data) self.assertEqual( container.digest_panels[spec.frequency].window_length, @@ -1052,7 +1052,7 @@ class TestHistoryContainerResize(TestCase): data_frequency=data_frequency, ) - container.ensure_spec(new_spec, initial_dt) + container.ensure_spec(new_spec, initial_dt, bar_data) if bar_count > 1: digest_panel = container.digest_panels[new_spec.frequency] @@ -1109,7 +1109,7 @@ class TestHistoryContainerResize(TestCase): data_frequency=data_frequency, ) - container.ensure_spec(new_spec, initial_dt) + container.ensure_spec(new_spec, initial_dt, bar_data) if bar_count > 1: digest_panel = container.digest_panels[new_spec.frequency] diff --git a/tests/utils/test_events.py b/tests/utils/test_events.py index 303647b0..fed4818b 100644 --- a/tests/utils/test_events.py +++ b/tests/utils/test_events.py @@ -15,6 +15,7 @@ import datetime import random from itertools import islice +from six import iteritems from six.moves import range, map from nose_parameterized import parameterized from unittest import TestCase @@ -216,7 +217,7 @@ class RuleTestCase(TestCase): return # This is the base class testing, it is always complete. dem = { - k for k, v in vars(zipline.utils.events).iteritems() + k for k, v in iteritems(vars(zipline.utils.events)) if isinstance(v, type) and issubclass(v, self.class_) and v is not self.class_ diff --git a/zipline/algorithm.py b/zipline/algorithm.py index a9438e02..a3a3c504 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -238,6 +238,8 @@ class TradingAlgorithm(object): if 'data_frequency' in kwargs: self.data_frequency = kwargs.pop('data_frequency') + self._most_recent_data = None + # Subclasses that override initialize should only worry about # setting self.initialized = True if AUTO_INITIALIZE is # is manually set to False. @@ -261,6 +263,7 @@ class TradingAlgorithm(object): self._before_trading_start(self) def handle_data(self, data): + self._most_recent_data = data if self.history_container: self.history_container.update(data, self.datetime) @@ -924,12 +927,13 @@ class TradingAlgorithm(object): self.history_specs[history_spec.key_str] = history_spec if self.initialized: if self.history_container: - self.history_container.ensure_spec(history_spec, self.datetime) + self.history_container.ensure_spec( + history_spec, self.datetime, self._most_recent_data, + ) else: self.history_container = self.history_container_class( - self.trade_sources.history_backfill, self.history_specs, - self.multiverse.current_sids, + self.current_universe(), self.sim_params.first_open, self.sim_params.data_frequency, ) @@ -952,9 +956,11 @@ class TradingAlgorithm(object): self.current_universe(), self.datetime, self.sim_params.data_frequency, - shift_digest=True, + bar_data=self._most_recent_data, ) - self.history_container.ensure_spec(spec, self.datetime) + self.history_container.ensure_spec( + spec, self.datetime, self._most_recent_data, + ) return self.history_specs[spec_key] @api_method diff --git a/zipline/history/history_container.py b/zipline/history/history_container.py index 84498669..29e06d25 100644 --- a/zipline/history/history_container.py +++ b/zipline/history/history_container.py @@ -183,7 +183,7 @@ class HistoryContainer(object): initial_sids, initial_dt, data_frequency, - shift_digest=False): + bar_data=None): """ A container to hold a rolling window of historical data within a user's algorithm. @@ -196,10 +196,9 @@ class HistoryContainer(object): initial_dt (datetime): The datetime to start collecting history from. - shift_digest (bool): If True, then the digest panels will be created - shifted back by one bar, this is to facilitate the creation of a - HistoryContainer during a call to handle_data within - TradingAlgorithm. This is False by default. + bar_data (BarData): If this container is being constructed during + handle_data, this is the BarData for the current bar to fill the + buffer with. If this is constructed elsewhere, it is None. Returns: An instance of a new HistoryContainer @@ -227,11 +226,11 @@ class HistoryContainer(object): # completed. When a frequency period rolls over, these minutes are # digested using some sort of aggregation call on the panel (e.g. `sum` # for volume, `max` for high, `min` for low, etc.). - self.buffer_panel = self.create_buffer_panel(initial_dt) + self.buffer_panel = self.create_buffer_panel(initial_dt, bar_data) # Dictionaries with Frequency objects as keys. self.digest_panels, self.cur_window_starts, self.cur_window_closes = \ - self.create_digest_panels(initial_sids, initial_dt, shift_digest) + self.create_digest_panels(initial_sids, initial_dt, bar_data) # Helps prop up the prior day panel against having a nan, when the data # has been seen. @@ -287,7 +286,7 @@ class HistoryContainer(object): return iterkeys(self.largest_specs) @with_environment() - def _add_frequency(self, spec, dt, env=None): + def _add_frequency(self, spec, dt, data, env=None): """ Adds a new frequency to the container. This reshapes the buffer_panel if needed. @@ -304,7 +303,7 @@ class HistoryContainer(object): # If the data_frequencies are not the same, then we need to # create a fresh buffer. self.buffer_panel = self.create_buffer_panel( - dt, shift_digest=True, + dt, bar_data=data, ) new_buffer_len = None else: @@ -495,7 +494,7 @@ class HistoryContainer(object): return self._create_panel(dt, spec, env=env) - def ensure_spec(self, spec, dt): + def ensure_spec(self, spec, dt, bar_data): """ Ensure that this container has enough space to hold the data for the given spec. This returns a HistoryContainerDelta to represent the @@ -506,7 +505,9 @@ class HistoryContainer(object): if spec.field not in self.fields: updated['field'] = self._add_field(spec.field) if spec.frequency not in self.largest_specs: - updated['frequency_delta'] = self._add_frequency(spec, dt) + updated['frequency_delta'] = self._add_frequency( + spec, dt, bar_data, + ) if spec.bar_count > self.largest_specs[spec.frequency].bar_count: updated['length_delta'] = self._add_length(spec, dt) return HistoryContainerDelta(**updated) @@ -550,7 +551,7 @@ class HistoryContainer(object): def create_digest_panels(self, initial_sids, initial_dt, - shift_digest, + bar_data, env=None): """ Initialize a RollingPanel for each unique panel frequency being stored @@ -577,7 +578,7 @@ class HistoryContainer(object): continue dt = initial_dt - if shift_digest: + if bar_data is not None: dt = largest_spec.frequency.prev_bar(dt) rp = self._create_digest_panel( @@ -592,7 +593,7 @@ class HistoryContainer(object): return panels, first_window_starts, first_window_closes - def create_buffer_panel(self, initial_dt): + def create_buffer_panel(self, initial_dt, bar_data): """ Initialize a RollingPanel containing enough minutes to service all our frequencies. @@ -609,6 +610,11 @@ class HistoryContainer(object): initial_dt, spec, ) self.buffer_spec = spec + + if bar_data: + frame = self.frame_from_bardata(bar_data, initial_dt) + rp.add_frame(initial_dt, frame) + return rp def convert_columns(self, values):