better variable names and PreTransformLayer

zipline/gens/composites.py

@@ -1,17 +1,18 @@
-
+from zipline.gens.utils import roundrobin
 from zipline.gens.feed import FeedGen
-from zipline.gens.tradegen import SpecificEquityTrades
-from zipline.gens.transform 
 
 
+def PreTransformLayer(sources, source_ids):
+    """
+    A generator that takes a tuple of sources and a list ids, piping
+    their output into a feed_gen.
+    """
+    stream_in = roundrobin(*sources)
+    return FeedGen(stream_in, source_ids)
 
+def TransformLayer(feed_stream, tnfms):
+    """ """
+    pass
 
-def PreTransformLayer(sources):
-    """A generator that takes a list of sources and runs their output
-    through a FeedGen."""
-    not_finished = len #NOT DONE
-
-    while not_finished:
-        
-        
+    
     

zipline/gens/feed.py

@@ -13,7 +13,8 @@ from collections import deque, defaultdict
 
 from zipline import ndict
 from zipline.gens.utils import hash_args, assert_datasource_protocol, \
-    assert_trade_protocol, assert_datasource_unframe_protocol
+    assert_trade_protocol, assert_datasource_unframe_protocol, \
+    assert_feed_protocol
 
 import zipline.protocol as zp
 
@@ -49,7 +50,7 @@ def FeedGen(stream_in, source_ids):
 
         while full(sources) and not done(sources):
             message = pop_oldest(sources)
-            assert feed_protocol(message)
+            assert_feed_protocol(message)
             yield message
 
     # We should have only a done message left in each queue.    

zipline/gens/test_feed.py

@@ -13,8 +13,10 @@ from collections import deque
 from zipline import ndict
 from zipline.gens.feed import FeedGen, full, done, queue_is_full,queue_is_done,\
     pop_oldest
-from zipline.gens.utils import stringify_args, assert_datasource_protocol,\
-    assert_trade_protocol, date_gen, alternate
+from zipline.gens.utils import hash_args, assert_datasource_protocol,\
+    assert_trade_protocol, alternate
+from zipline.gens.tradegens import date_gen, SpecificEquityTrades
+from zipline.gens.composites import PreTransformLayer
 
 import zipline.protocol as zp
 
@@ -98,12 +100,11 @@ class FeedGenTestCase(TestCase):
         l = list(feed_gen)
         assert l == expected
         
-
     def test_single_source(self):
         source_ids = ['a']
         # 100 events, increasing by a minute at a time.
         type = zp.DATASOURCE_TYPE.TRADE
-        dates = list(date_gen(n = 1))
+        dates = list(date_gen(count = 100))
         dates.append("DONE")
         
         # [('a', date1, type), ('a', date2, type), ... ('a', "DONE", type)]
@@ -125,7 +126,7 @@ class FeedGenTestCase(TestCase):
 
         # Set up source 'a'. Outputs 20 events with 2 minute deltas.
         delta_a = timedelta(minutes = 2)
-        dates_a = list(date_gen(delta = delta_a, n = 20))
+        dates_a = list(date_gen(delta = delta_a, count = 20))
         dates_a.append("DONE")
 
         events_a_args = zip(cycle(['a']), iter(dates_a), cycle([type]))
@@ -133,7 +134,7 @@ class FeedGenTestCase(TestCase):
         
         # Set up source 'b'. Outputs 10 events with 1 minute deltas.
         delta_b = timedelta(minutes = 1)
-        dates_b = list(date_gen(delta = delta_b, n = 10))
+        dates_b = list(date_gen(delta = delta_b, count = 10))
         dates_b.append("DONE")
 
         events_b_args = zip(cycle(['b']), iter(dates_b), cycle([type]))
@@ -153,12 +154,41 @@ class FeedGenTestCase(TestCase):
         sequential = chain(iter(events_a), iter(events_b))
         self.run_FeedGen(sequential, expected, source_ids)
     
-        def test_with_specific_equity(self):
-            
-            
+    def test_full_feed_layer(self):
+        filter = [1,2]
         
+        source_a = SpecificEquityTrades(sids = [1,2,3,4],
+                                        start = datetime(2012,6,6,0),
+                                        delta = timedelta(minutes=1),
+                                        filter = filter
+        )
+        id_a = "SpecificEquityTradesd175237b28d2f52df208c97cf4af896e"
+
+        # Change the internal sid list to give us a different hash.
+        source_b = SpecificEquityTrades(sids = [1,2,3,5],
+                                        start = datetime(2012,6,6,0),
+                                        delta = timedelta(minutes=1),
+                                        filter = filter
+        )
         
+        id_b = 'SpecificEquityTrades2bf2c2d6d01d4dbfc0b2818438ea8151'
+
+        # Change the internal sid list to give us a different hash.
+        source_c = SpecificEquityTrades(sids = [1,2,3,6],
+                                        start = datetime(2012,6,6,0),
+                                        delta = timedelta(minutes=1),
+                                        filter = filter
+        )
+        id_c = 'SpecificEquityTrades16f7437db2d14e5373ef20025f49a3fe'
+
+        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
+    
 def mock_data_unframe(source_id, dt, type):
     event = ndict()
     event.source_id = source_id

zipline/gens/test_mongods.py

@@ -10,7 +10,7 @@ from itertools import izip, izip_longest
 from datetime import datetime, timedelta
 
 from zipline.gens.mongods import create_pymongo_iterator, MongoTradeHistoryGen
-from zipline.gens.utils import stringify_args, assert_datasource_protocol,\
+from zipline.gens.utils import hash_args, assert_datasource_protocol,\
     assert_trade_protocol, mock_raw_event
 
 import zipline.protocol as zp
@@ -107,7 +107,7 @@ class TestMongoDataGenerator(TestCase):
                 for field in iter(['sid', 'dt', 'price', 'volume']):
                     assert db[field] == expected[field]
                 
-                # Expected output of stringify_args:
+                # Expected output of hash_args:
                 assert db['source_id'] == \
                     'MongoTradeHistoryGen983a27fd0710414239a5cde71ef5a8fc'
                 

zipline/gens/tradegens.py

@@ -2,53 +2,76 @@ import random
 from itertools import chain, repeat, cycle, ifilter, izip
 from datetime import datetime, timedelta
 
-from zipline.utils.factory import create_trade, create_trade
-from zipline.gens.utils import date_gen
+from zipline.utils.factory import create_trade
+from zipline.gens.utils import hash_args
 
-def mock_prices(n, rand = False):
-    """Utility to generate a set of prices. By default
-    cycles through values from 0.0 to 10.0 n times.  Optional
-    flag to give random values between 0.0 and 10.0"""
+def date_gen(start = datetime(2012, 6, 6, 0),
+             delta = timedelta(minutes = 1), 
+             count = 100):
+    """
+    Utility to generate a stream of dates.
+    """
+    return (start + (i * delta) for i in xrange(count))
+
+def mock_prices(count, rand = False):
+    """
+    Utility to generate a stream of mock prices. By default
+    cycles through values from 0.0 to 10.0, n times.  Optional
+    flag to give random values between 0.0 and 10.0
+    """
 
     if rand:
-        return (random.uniform(0.0, 10.0) for i in xrange(n))
+        return (random.uniform(0.0, 10.0) for i in xrange(count))
     else:
-        return (float(i % 11) for i in xrange(1,n+1))
+        return (float(i % 11) for i in xrange(1,count+1))
 
-def mock_volumes(n, rand = False):
-    """Utility to generate a set of volumes. By default cycles
+def mock_volumes(count, rand = False):
+    """
+    Utility to generate a set of volumes. By default cycles
     through values from 100 to 1000, incrementing by 50.  Optional
-    flag to give random values between 100 and 1000. """
-
+    flag to give random values between 100 and 1000. 
+    """
     if rand:
-        return (random.randrange(100, 1000) for i in xrange(n))
+        return (random.randrange(100, 1000) for i in xrange(count))
     else:
-        return ((i * 50)%900 + 100 for i in xrange(n))
+        return ((i * 50)%900 + 100 for i in xrange(count))
     
 def fuzzy_dates(count = 500):
-    """Add +-10 seconds to each event from a date_gen.  Note that
-    this still guarantees sorting, since the default is minute separation
-    of events."""
-    for date in date_gen(n = count):
+    """
+    Add +-10 seconds to each event from a date_gen.  Note that this
+    still guarantees sorting, since the default on date_gen is minute
+    separation of events.
+    """
+    for date in date_gen(count = count):
         yield date + timedelta(seconds = random.randint(-10, 10)) 
 
-def SpecificEquityTrades(count = 500, sids = [1, 2], event_list = None, filter = None):
-                             
-    """Returns the first n events of event_list if specified. 
-    Otherwise generates a sensible stream of events."""
-    
+def SpecificEquityTrades(count = 500, 
+                         sids = [1, 2],
+                         start = datetime(2012, 6, 6, 0),
+                         delta = timedelta(minutes = 1),
+                         event_list = None, 
+                         filter = None):
+    """
+    Yields all events in event_list that match the given sid_filter.
+    If no event_list is specified, generates an internal stream of events
+    to filter.  Returns all events if filter is None.
+    """
+    arg_string = hash_args(count, sids, start, delta, filter)
+    namestring = "SpecificEquityTrades" + arg_string
+
     if event_list:
         unfiltered = (event for event in event_list)
     
     else:
-        dates = date_gen(n = count)
+        dates = date_gen(count = count, start = start, delta = delta)
         prices = mock_prices(count)
         volumes = mock_volumes(count)
         sids = cycle(iter(sids))
         
         arg_gen = izip(sids, prices, volumes, dates)
         
-        unfiltered = (create_trade(*args) for args in arg_gen)
+        unfiltered = (create_trade(*args, source_id = namestring)
+                      for args in arg_gen)
     if filter:
         filtered = ifilter(lambda event: event.sid in filter, unfiltered)
     else:
@@ -72,9 +95,6 @@ def RandomEquityTrades(count = 500, sids = [1,2], filter = None):
         filtered = unfiltered
     return filtered
 
-if __name__ == "__main__":
-    rand = RandomEquityTrades()
-    pass
-#     x = mock_volumes(500)
+# if __name__ == "__main__":
 #     import nose.tools; nose.tools.set_trace()
 #     trades = SpecificEquityTrades(filter = [1])

zipline/gens/utils.py

@@ -16,16 +16,25 @@ def mock_raw_event(sid, dt):
     }
     return event
 
-def date_gen(start = datetime(2012, 6, 6, 0), delta = timedelta(minutes = 1), n = 100):
-    return (start + (i * delta) for i in xrange(n))
-
 def alternate(g1, g2):
+    """Specialized version of roundrobin for just 2 generators."""
     for e1, e2 in izip_longest(g1, g2):
         if e1 != None:
             yield e1
         if e2 != None:
             yield e2
 
+def roundrobin(*args):
+    """
+    Takes N generators, pulling one element off each until all inputs
+    are empty.  
+    """
+    for elem_tuple in izip_longest(*args):
+        for value in elem_tuple:
+            if value != None:
+                yield value
+
+
 def hash_args(*args, **kwargs):
     """Define a unique string for any set of representable args."""
     arg_string = '_'.join([str(arg) for arg in args])

zipline/utils/factory.py

@@ -69,9 +69,9 @@ def create_trading_environment(year=2006):
 
     return trading_environment
 
-def create_trade(sid, price, amount, datetime):
+def create_trade(sid, price, amount, datetime, source_id = "test_factory"):
     row = zp.ndict({
-        'source_id' : "test_factory",
+        'source_id' : source_id,
         'type'      : zp.DATASOURCE_TYPE.TRADE,
         'sid'       : sid,
         'dt'        : datetime,