diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index 10e3cd2e..3af77fc1 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -1,5 +1,5 @@ import datetime -from itertools import tee, starmap +from itertools import tee, starmap, chain from collections import namedtuple from zipline.gens.tradegens import SpecificEquityTrades @@ -66,7 +66,7 @@ def merged_transforms(sorted_stream, *transforms): # Pipe the stream into merge. merged = merge(to_merge, namestrings) # Return the merged events. - return merged + return add_done(dt_aliased) def sequential_transforms(stream_in, *transforms): """ @@ -87,7 +87,7 @@ def sequential_transforms(stream_in, *transforms): stream_in) dt_aliased = alias_dt(stream_out) - return dt_aliased + return add_done(dt_aliased) def alias_dt(stream_in): """ @@ -95,10 +95,11 @@ def alias_dt(stream_in): """ for message in stream_in: message['datetime'] = message['dt'] - yield message + yield message - - +# Add a done message to a stream. +def add_done(stream_in): + return chain(stream_in, [done_message('Composite')]) diff --git a/zipline/gens/merge.py b/zipline/gens/merge.py index 0689d507..09e1f943 100644 --- a/zipline/gens/merge.py +++ b/zipline/gens/merge.py @@ -51,8 +51,7 @@ def merge(stream_in, tnfm_ids): assert len(queue) == 1, "Bad queue in merge on exit: %s" % queue assert queue[0].dt == "DONE", \ "Bad last message in merge on exit: %s" % queue - yield done_message('Merge') - + def merge_one(sources): event_fields = ndict() diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 50326368..6f3323bd 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -97,7 +97,7 @@ class TradeSimulationClient(object): ordering_client.state, self.algo, ) - + # The algorithm will yield a daily_results message (as # calculated by the performance tracker) at the end of each # day. It will also yield a risk report at the end of the @@ -105,7 +105,6 @@ class TradeSimulationClient(object): for message in self.algo_sim: yield message - class AlgorithmSimulator(object): def __init__(self, stream_in, order_book, algo): @@ -215,16 +214,18 @@ class AlgorithmSimulator(object): for event in self.stream_in: # Yield any perf messages received to be relayed back to # the browser. + if event.perf_message: yield event.perf_message del event['perf_message'] - if event.dt == "DONE": - if self.this_snapshot_dt: - # stop iteration happened - # mid-snapshot, so we have a universe - # snapshot that is not yet processed - # by the algorithm. - self.simulate_current_snapshot() + + if event.dt == "DONE": + if self.this_snapshot_dt: + # stop iteration happened + # mid-snapshot, so we have a universe + # snapshot that is not yet processed + # by the algorithm. + self.simulate_current_snapshot() break # This should only happen for the first event we run.