added done message to SpecificEquity

zipline/gens/feed.py

@@ -52,7 +52,8 @@ def FeedGen(stream_in, source_ids):
             message = pop_oldest(sources)
             assert_feed_protocol(message)
             yield message
-
+    
+    import nose.tools; nose.tools.set_trace()
     # We should have only a done message left in each queue.    
     for queue in sources.itervalues():
         assert len(queue) == 1, "Bad queue in FeedGen on exit: %s" % queue

zipline/gens/test_feed.py

@@ -184,10 +184,8 @@ class FeedGenTestCase(TestCase):
         sources = (source_a, source_b, source_c)
         source_ids = [id_a, id_b, id_c]
         
-        import nose.tools; nose.tools.set_trace()
         feed_out = PreTransformLayer(sources, source_ids)
-        for i in feed_out:
-            print i
+        l = list(feed_out)
     
 def mock_data_unframe(source_id, dt, type):
     event = ndict()

zipline/gens/tradegens.py

@@ -3,7 +3,7 @@ from itertools import chain, repeat, cycle, ifilter, izip
 from datetime import datetime, timedelta
 
 from zipline.utils.factory import create_trade
-from zipline.gens.utils import hash_args
+from zipline.gens.utils import hash_args, mock_done
 
 def date_gen(start = datetime(2012, 6, 6, 0),
              delta = timedelta(minutes = 1), 
@@ -77,7 +77,9 @@ def SpecificEquityTrades(count = 500,
     else:
         filtered = unfiltered
 
-    return filtered
+    # Add a done message to the end of the stream.
+    out = chain(filtered, iter([mock_done(namestring)]))
+    return out    
 
 def RandomEquityTrades(count = 500, sids = [1,2], filter = None):
     dates = fuzzy_dates(500)

zipline/gens/utils.py

@@ -16,6 +16,9 @@ def mock_raw_event(sid, dt):
     }
     return event
 
+def mock_done(source_id):
+    return ndict({'dt': "DONE", "source_id" : source_id, 'type' : 0})
+              
 def alternate(g1, g2):
     """Specialized version of roundrobin for just 2 generators."""
     for e1, e2 in izip_longest(g1, g2):