full sequencing system (minus done from xforms)

zipline/gens/composites.py

@@ -1,9 +1,11 @@
+import datetime
 from itertools import tee
 
 from zipline.gens.tradegens import SpecificEquityTrades
 from zipline.gens.utils import roundrobin, hash_args
 from zipline.gens.feed import FeedGen
 from zipline.gens.merge import MergeGen
+from zipline.gens.transform import StatefulTransformGen
 
 def PreTransformLayer(sources, source_args, source_kwargs):
     """
@@ -14,15 +16,17 @@ def PreTransformLayer(sources, source_args, source_kwargs):
     assert len(sources) == len(source_args) == len(source_kwargs)
     # Package up sources and arguments.
     arg_bundles = zip(sources, source_args, source_kwargs)
+    
     # Calculate namestring hashes to pass to FeedGen.
     namestrings = [source.__name__ + hash_args(*args, **kwargs)
                    for source, args, kwargs in arg_bundles]
     # Pass each source its arguments.
     initialized = tuple(source(*args, **kwargs)
                         for source, args, kwargs in arg_bundles)
-
+    
     stream_in = roundrobin(*initialized)
-    return FeedGen(stream_in, source_ids)
+    return FeedGen(stream_in, namestrings)
+
 
 def TransformLayer(feed_stream, tnfms, tnfm_args, tnfm_kwargs):
     """ 
@@ -34,6 +38,7 @@ def TransformLayer(feed_stream, tnfms, tnfm_args, tnfm_kwargs):
     tnfm_kwargs should be a list of dictionaries representing keyword
     arguments to each transform.
     """
+
     # We should have as many sets of args as we have transforms.
     assert len(tnfms) == len(tnfm_args) == len(tnfm_kwargs)
 
@@ -42,21 +47,34 @@ def TransformLayer(feed_stream, tnfms, tnfm_args, tnfm_kwargs):
     
     # Package each stream copy with a transform and set of args.  Use a list
     # so that we can re-use this for calculating hashes.
-    bundles = zip(split, iter(tnfms), iter(tnfm_args), iter(tnfm_kwargs))
-    
-    # Convert the argument bundles into a tuple of transform objects.
-    transformed = tuple((tnfm(stream, *args, **kwargs)
-                   for stream, tnfm, args, kwargs in iter(bundles)))
+    bundles = zip(split, tnfms, tnfm_args, tnfm_kwargs)
+ 
+    tnfm_gens = [StatefulTransformGen(stream, tnfm, *args, **kwargs)
+                 for stream, tnfm, args, kwargs in bundles]
     
+    # Generate expected hashes for each transform
+    hashes = [tnfm.__name__ + hash_args(*args, **kwargs) 
+              for _, tnfm, args, kwargs in bundles]
+              
     # Roundrobin the outputs of our transforms to create a single flat stream.
-    to_merge = roundrobin(*transformed)
-    
-    merged = MergeGen()
-    
+    to_merge = roundrobin(*tnfm_gens)
     
+    # Pipe the stream into MergeGen.
+    merged = MergeGen(to_merge, hashes)
+    return merged
+
 if __name__ == "__main__":
+    from zipline.gens.transform import MovingAverage, Passthrough
     
-    source = SpecificEquityTrades()
+    import nose.tools; nose.tools.set_trace()
+    source = SpecificEquityTrades
+    feed_out = PreTransformLayer((source,), ((),), ({},))
+
+    transforms = [MovingAverage, Passthrough]
+    args = [(datetime.timedelta(days = 1), ['price']), ()]
+    kwargs = [{}, {}]
+    
+    tlayer = TransformLayer(feed_out, transforms, args, kwargs)
     
     
     

zipline/gens/feed.py

@@ -39,7 +39,7 @@ def FeedGen(stream_in, source_ids):
         # Incoming messages should be the output of DATASOURCE_UNFRAME.
         assert_datasource_unframe_protocol(message), \
             "Bad message in FeedGen: %s" % message
-
+        
         # Only allow messages from sources we expect.
         assert message.source_id in sources, "Unexpected source: %s" % message
 
@@ -61,8 +61,9 @@ def FeedGen(stream_in, source_ids):
 
 def full(sources):
     """
-    Feed is full when every internal queue has at least one message. Note that
-    this include DONE messages, so done(sources) is True only if full(sources).
+    Feed is full when every internal queue has at least one
+    message. Note that this include DONE messages, so done(sources) is
+    True only if full(sources).
     """
     assert isinstance(sources, dict)
     return all( (queue_is_full(source) for source in sources.itervalues()) )

zipline/gens/merge.py

@@ -13,7 +13,7 @@ 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_merge_protocol
 
 import zipline.protocol as zp
 
@@ -25,44 +25,46 @@ def MergeGen(stream_in, tnfm_ids):
     and merge them together into an event.  We raise an error if we 
     do not receive the same number of events from all sources.
     """
-    
-    assert isinstance(source_ids, list)
 
+    assert isinstance(tnfm_ids, list)
+    
     # Set up an internal queue for each expected source.
-    sources = {}
-    for id in source_ids:
-        assert isinstance(id, basestring), "Bad source_id %s" % source_id
-        sources[id] = deque()
+    tnfms = {}
+    for id in tnfm_ids:
+        assert isinstance(id, basestring), "Bad source_id %s" % id
+        tnfms[id] = deque()
 
     # Process incoming streams.
     for message in stream_in:
-        assert isinstance(message, ndict), \
+        assert isinstance(message, tuple), \
             "Bad message in MergeGen: %s" %message
-        assert message.tnfm_id in tnfm_ids, \
-            "Message from unexpected tnfm: %s, %s" % (message, tnfm_ids)
+        assert len(message) == 2
+        id, value = message
+        assert id in tnfm_ids, \
+            "Message from unexpected tnfm: %s, %s" % (id, tnfm_ids)
+        assert isinstance(value, ndict), "Bad message in MergeGen: %s" %message
 
-        assert message.has_key('value')
-
-        source[message.tnfm_id].append(message)
+        tnfms[id].append(value)
 
         # Only pop messages when we have a pending message from
         # all datasources. Stop if all sources have signalled done.
 
-        while full(sources) and not done(sources):
-            message = merge_one(sources)
-            assert merge_protocol(message)
+        while full(tnfms) and not done(tnfms):
+            message = merge_one(tnfms)
+            assert_merge_protocol(tnfm_ids, message)
             yield message
 
     # We should have only a done message left in each queue.    
-    for queue in sources.itervalues():
+    for queue in tnfms.itervalues():
         assert len(queue) == 1, "Bad queue in MergeGen on exit: %s" % queue
         assert queue[0].dt == "DONE", \
             "Bad last message in MergeGen on exit: %s" % queue
 
 def merge_one(sources):
     output = ndict()
-    for queue in sources.itervalues():
-        output.merge(queue.popleft())
+    for key, queue in sources.iteritems():
+        new_xform = ndict({key: queue.popleft()})
+        output.merge(new_xform)
     return output
 
 

zipline/gens/test_feed.py

@@ -155,38 +155,61 @@ class FeedGenTestCase(TestCase):
         self.run_FeedGen(sequential, expected, source_ids)
 
     def test_full_feed_layer(self):
+        
         filter = [1,2]
-        #Set up source a.
+        #Set up source a. One hour between events.
         args_a = tuple()
         kwargs_a = {'sids'   : [1,2,3,4],
                     'start'  : datetime(2012,6,6,0),
-                    'delta'  : timedelta(minutes = 1),
+                    'delta'  : timedelta(hours = 1),
                     'filter' : filter
         }
-        #Set up source b.        
+        #Set up source b. One day between events.       
         args_b = tuple()
-        kwargs_b = {'sids'   : [1,2,3,5],
+        kwargs_b = {'sids'   : [1,2,3,4],
                     'start'  : datetime(2012,6,6,0),
-                    'delta'  : timedelta(minutes = 1),
+                    'delta'  : timedelta(days = 1),
                     'filter' : filter
         }
-        #Set up source c.
+        #Set up source c. One minute between events.
         args_c = tuple()
-        kwargs_c = {'sids'   : [1,2,3,5],
+        kwargs_c = {'sids'   : [1,2,3,4],
+                    'start'  : datetime(2012,6,6,0),
+                    'delta'  : timedelta(minutes = 1),
+                    'filter' : filter
+        }
+        # Set up source d. This should produce no events because the
+        # internal sids don't match the filter.
+        args_d = tuple()
+        kwargs_d = {'sids'   : [3,4],
                     'start'  : datetime(2012,6,6,0),
                     'delta'  : timedelta(minutes = 1),
                     'filter' : filter
         }
         
-        sources = tuple(SpecificEquityTrades) * 3
-        source_args = (args_a, args_b, args_c)
-        source_kwargs = (kwargs_a, kwargs_b, kwargs_c)
+        sources = (SpecificEquityTrades,) * 4
+        source_args = (args_a, args_b, args_c, args_d)
+        source_kwargs = (kwargs_a, kwargs_b, kwargs_c, kwargs_d)
         
+        # Generate our expected source_ids.
+        zip_args = zip(source_args, source_kwargs)
+        expected_ids = ["SpecificEquityTrades" + hash_args(*args, **kwargs)
+                        for args, kwargs in zip_args]
+        
+        # Pipe our sources into feed.
         feed_out = PreTransformLayer(sources, source_args, source_kwargs)
+        
+        # Read all the values from feed and assert that they arrive in
+        # the correct sorting with the expected hash values.
         to_list = list(feed_out)
         copy = to_list[:]
+        for e in to_list:
+            # All events should match one of our expected source_ids.
+            assert e.source_id in expected_ids
+            # But none of them should match source_d.
+            assert e.source_id != hash_args(*args_d, **kwargs_d)
+
         expected = sorted(copy, compare_by_dt_source_id)
-        
         assert to_list == expected
     
 def mock_data_unframe(source_id, dt, type):

zipline/gens/tradegens.py

@@ -45,47 +45,71 @@ def fuzzy_dates(count = 500):
     for date in date_gen(count = count):
         yield date + timedelta(seconds = random.randint(-10, 10)) 
 
-def SpecificEquityTrades(count = 500, 
-                         sids = [1, 2],
-                         start = datetime(2012, 6, 6, 0),
-                         delta = timedelta(minutes = 1),
-                         event_list = None, 
-                         filter = None):
+def SpecificEquityTrades(*args, **config):
     """
     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
+    # We shouldn't get any positional arguments.
+    assert args == ()
 
+    # Unpack config dictionary with default values.
+    count = config.get('count', 500)
+    sids = config.get('sids', [1, 2])
+    start = config.get('start', datetime(2012, 6, 6, 0))
+    delta = config.get('delta', timedelta(minutes = 1))
+
+    # Default to None for event_list and filter.
+    event_list = config.get('event_list') 
+    filter = config.get('filter')
+
+    arg_string = hash_args(*args, **config)
+    namestring = "SpecificEquityTrades" + arg_string
+    # If we have an event_list, ignore the other arguments and use the list.
+    # TODO: still append our namestring?
     if event_list:
         unfiltered = (event for event in event_list)
-    
+
+    # Set up iterators for each expected field.
     else:
         dates = date_gen(count = count, start = start, delta = delta)
         prices = mock_prices(count)
         volumes = mock_volumes(count)
-        sids = cycle(iter(sids))
-        
+        sids = cycle(sids)
+
+        # Combine the iterators into a single iterator of arguments
         arg_gen = izip(sids, prices, volumes, dates)
-        
+
+        # Convert argument packages into events.
         unfiltered = (create_trade(*args, source_id = namestring)
                       for args in arg_gen)
+
+    # If we specified a sid filter, filter out elements that don't match the filter.
     if filter:
         filtered = ifilter(lambda event: event.sid in filter, unfiltered)
+
+    # Otherwise just use all events.
     else:
         filtered = unfiltered
 
-    # Add a done message to the end of the stream.
-    out = chain(filtered, iter([mock_done(namestring)]))
-    return out    
+    # Add a done message to the end of the stream. For a live
+    # datasource this would be handled by the containing Component.
+    out = chain(filtered, [mock_done(namestring)])
+    return out
 
-def RandomEquityTrades(count = 500, sids = [1,2], filter = None):
-    dates = fuzzy_dates(500)
-    prices = mock_prices(500, rand = True)
-    volumes = mock_volumes(500, rand = True)
-    sids = cycle(iter(sids))
+def RandomEquityTrades(*args, **config):
+    # We shouldn't get any positional args.
+    assert args == ()
+    
+    count = config.get('count', 500)
+    sids = config.get('sids', [1,2])
+    filter = config.get('filter')
+
+    dates = fuzzy_dates(count)
+    prices = mock_prices(count, rand = True)
+    volumes = mock_volumes(count, rand = True)
+    sids = cycle(sids)
 
     arg_gen = izip(sids, prices, volumes, dates)
     

zipline/gens/transform.py

@@ -20,21 +20,18 @@ from zipline.gens.utils import assert_feed_unframe_protocol, \
 
 import zipline.protocol as zp
 
-def PassthroughTransformGen(stream_in):
-    """Trivial transform for event forwarding."""
+class Passthrough(object):
+    """
+    Trivial function for forwarding events.
+    """
+    def __init__(self):
+        pass
 
-    # hash_args with no arguments is the same as:
-    # hasher = hashlib.md5() 
-    # hasher.update(":"); 
-    # hashlib.md5.digest().
-
-    namestring = "Passthrough" + hash_args()
-
-    for message in stream_in:
-        assert_feed_unframe_protocol(message)
-        out_value = message
-        assert_transform_protocol(out_value)
-        yield (namestring, out_value)
+    def update(self, event):        
+        assert isinstance(event, ndict),"Bad event in Passthrough: %s" % event
+        assert event.has_key('sid'), "No sid in Passthrough: %s" % event
+        assert event.has_key('dt'), "No dt in Passthorughz: %s" % event
+        return event
 
 def FunctionalTransformGen(stream_in, fun, *args, **kwargs):
     """
@@ -75,16 +72,6 @@ def StatefulTransformGen(stream_in, tnfm_class, *args, **kwargs):
         assert_transform_protocol(out_value)
         yield (namestring, out_value)
 
-def MovingAverageTransformGen(stream_in, days, fields):
-    """
-    Generator that uses the MovingAverage state class to calculate
-    a moving average for all stocks over a specified number of days.
-    """
-    return StatefulTransformGen(stream_in, 
-                                MovingAverage, 
-                                timedelta(days=days), 
-                                fields)
-
 class MovingAverage(object):
     """
     Class that maintains a dictionary from sids to EventWindows
@@ -112,8 +99,9 @@ class MovingAverage(object):
         
         assert isinstance(event, ndict),"Bad event in MovingAverage: %s" % event
         assert event.has_key('sid'), "No sid in MovingAverage: %s" % event
+        assert event.has_key('dt'), "No dt in MovingAverage: %s" % event
         
-        output = ndict({'sid': event.sid})
+        output = ndict({'sid': event.sid, 'dt': event.dt})
         # This will create a new EventWindow if this is the first
         # message for this sid.
         window = self.sid_windows[event.sid]

zipline/gens/utils.py

@@ -96,3 +96,8 @@ def assert_transform_protocol(event):
     """Transforms should return an ndict to be merged by MergeGen."""
     assert isinstance(event, ndict)
 
+def assert_merge_protocol(tnfm_ids, message):
+    """Merge should output an ndict with a field for each id in its transform set."""
+    assert isinstance(message, ndict)
+    assert set(tnfm_ids) == set(message.keys())
+

zipline/gens/zmq_gens.py

@@ -0,0 +1,16 @@
+import zmq
+
+import zipline.protocol as zp
+
+def gen_from_zmq(poller, unframe):
+    """
+    A generator that takes an initialized zmq poller and yields
+    messages from the poller until it gets a zp.CONTROL_PROTOCOL.DONE.
+    """
+    while True:
+        message = poller.recv()
+        if message = zp.CONTROL_PROTOCOL.DONE:
+            yield "DONE"
+            break
+        else:
+            yield unframe(message)