diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py index 97b338d2..4c0e2919 100644 --- a/tests/test_perf_tracking.py +++ b/tests/test_perf_tracking.py @@ -1060,28 +1060,34 @@ class TestPerformanceTracker(unittest.TestCase): bar_event_2 = factory.create_trade( 'bar', 11.0, 20, start_dt + datetime.timedelta(minutes=1)) - foo_event_3 = factory.create_trade( - 'foo', 12.0, 30, start_dt + datetime.timedelta(minutes=2)) + events = [ + foo_event_1, + bar_event_1, + foo_event_2, + bar_event_2 + ] - tracker.process_event(foo_event_1) - tracker.process_event(bar_event_1) - messages = tracker.process_event(foo_event_2) - tracker.process_event(bar_event_2) - messages += tracker.process_event(foo_event_3) + import operator + messages = {date: snapshot[0].perf_messages[0] for date, snapshot in + tracker.transform( + itertools.groupby( + events, + operator.attrgetter('dt')))} self.assertEquals(2, len(messages)) - self.assertEquals(1, len(messages[0]['intraday_perf']['transactions']), + msg_1 = messages[foo_event_1.dt] + msg_2 = messages[foo_event_2.dt] + + self.assertEquals(1, len(msg_1['intraday_perf']['transactions']), "The first message should contain one transaction.") # Check that transactions aren't emitted for previous events. - self.assertEquals(0, len(messages[1]['intraday_perf']['transactions']), + self.assertEquals(0, len(msg_2['intraday_perf']['transactions']), "The second message should have no transactions.") # Ensure that period_close moves through time. # Also, ensure that the period_closes are the expected dts. self.assertEquals(foo_event_1.dt, - messages[0]['intraday_perf']['period_close'], - messages[0]) + msg_1['intraday_perf']['period_close']) self.assertEquals(foo_event_2.dt, - messages[1]['intraday_perf']['period_close'], - messages[1]) + msg_2['intraday_perf']['period_close']) diff --git a/zipline/finance/performance.py b/zipline/finance/performance.py index 92d0bef8..06ddcd77 100644 --- a/zipline/finance/performance.py +++ b/zipline/finance/performance.py @@ -130,7 +130,6 @@ omitted). """ -import itertools import logbook import math @@ -206,24 +205,26 @@ class PerformanceTracker(object): """ Main generator work loop. """ - # Set the simulation date to be the first event we see. - peek_date, peek_snapshot = next(stream_in) - self.saved_dt = peek_date - - # Stitch back together the generator by placing the peeked - # event back in front - stream = itertools.chain([(peek_date, peek_snapshot)], - stream_in) - - for date, snapshot in stream: + for date, snapshot in stream_in: new_snapshot = [] - for event in snapshot: - messages = self.process_event(event) - if messages is not None: - event.perf_messages = messages - event.portfolio = self.get_portfolio() + if self.emission_rate == 'daily': + for event in snapshot: + messages = self.process_event(event) + if messages is not None: + event.perf_messages = messages + event.portfolio = self.get_portfolio() + new_snapshot.append(event) + + elif self.emission_rate == 'minute': + self.saved_dt = date + self.todays_performance.period_close = self.saved_dt + + for event in snapshot: + self.process_event(event) + event.perf_messages = [self.to_dict()] + event.portfolio = self.get_portfolio() new_snapshot.append(event) if new_snapshot: @@ -277,10 +278,7 @@ class PerformanceTracker(object): event.dt < self.last_close): messages.append(self.handle_market_close()) elif self.emission_rate == 'minute': - if event.dt > self.saved_dt: - self.todays_performance.period_close = self.saved_dt - messages.append(self.to_dict()) - self.saved_dt = event.dt + messages.append(self.to_dict()) if event.TRANSACTION: self.txn_count += 1 diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 6a58866f..23072506 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -285,10 +285,11 @@ class AlgorithmSimulator(object): perf_messages, risk_message = \ self.perf_tracker.handle_simulation_end() - for message in perf_messages: - message[self.perf_key]['recorded_vars'] =\ - self.algo.recorded_vars - yield message + if self.perf_tracker.emission_rate == 'daily': + for message in perf_messages: + message[self.perf_key]['recorded_vars'] =\ + self.algo.recorded_vars + yield message # When emitting minutely, it is still useful to have a final # packet with the entire days performance rolled up. @@ -298,7 +299,6 @@ class AlgorithmSimulator(object): ) daily_rollup['daily_perf']['recorded_vars'] = \ self.algo.recorded_vars - log.info("emitting daily rollup: %s" % daily_rollup) yield daily_rollup yield risk_message