done message for seq transforms

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')])
     
     
     

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()

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.