commit for fawce

zipline/gens/composites.py

@@ -0,0 +1,17 @@
+
+from zipline.gens.feed import FeedGen
+from zipline.gens.tradegen import SpecificEquityTrades
+from zipline.gens.transform 
+
+
+
+
+def PreTransformLayer(sources):
+    """A generator that takes a list of sources and runs their output
+    through a FeedGen."""
+    not_finished = len_
+
+    while not_finished:
+        
+        
+    

zipline/gens/merge.py

@@ -0,0 +1,99 @@
+"""
+Generator version of Merge.
+"""
+
+import pytz
+import logbook
+import pymongo
+import types
+
+from pymongo import ASCENDING
+from datetime import datetime, timedelta
+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
+
+import zipline.protocol as zp
+
+def MergeGen(stream_in, tnfm_ids):
+    """
+    A generator that takes a generator and a list of source_ids. We
+    maintain an internal queue for each id in source_ids. Once we 
+    have a message from every queue, we pop an event from each queue
+    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)
+
+    # 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()
+
+    # Process incoming streams.
+    for message in stream_in:
+        assert isinstance(message, ndict), \
+            "Bad message in MergeGen: %s" %message
+        assert message.tnfm_id in tnfm_ids, \
+            "Message from unexpected tnfm: %s, %s" % (message, tnfm_ids)
+
+        assert message.has_key('value')
+
+        source[message.tnfm_id].append(message)
+
+        # 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)
+            yield message
+
+    # We should have only a done message left in each queue.    
+    for queue in sources.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())
+    return output
+
+
+#TODO: This is replicated in feed.  Probably should be one source file.
+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).
+    """
+    assert isinstance(sources, dict)
+    return all( (queue_is_full(source) for source in sources.itervalues()) )
+
+def queue_is_full(queue):
+    assert isinstance(queue, deque) 
+    return len(queue) > 0
+
+def done(sources):
+    """Feed is done when all internal queues have only a "DONE" message."""
+    assert isinstance(sources, dict) 
+    return all( (queue_is_done(source) for source in sources.itervalues()) )
+
+def queue_is_done(queue):
+    assert isinstance(queue, deque)
+    if len(queue) == 0:
+        return False
+    if queue[0].dt == "DONE":
+        assert len(queue) == 1, "Message after DONE in FeedGen: %s" % queue
+        return True
+    else:
+        return False
+
+        
+        
+    

zipline/gens/test_feed.py

@@ -47,7 +47,7 @@ class FeedHelperTestCase(TestCase):
         assert queue_is_full(queue)
         assert queue_is_done(queue)
         
-    def test_sources_logic(self):
+    def test_pop_logic(self):
         sources = {}
         ids = ['a', 'b', 'c']
         for id in ids:
@@ -85,7 +85,6 @@ class FeedGenTestCase(TestCase):
     
     def setUp(self):
         pass
-        
     
     def tearDown(self):
         pass
@@ -154,15 +153,10 @@ 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_FeedGen_consistency(self):
-        
-#         source_ids = ['a', 'b']
-#         multiplied = source_ids * 5
-#         perms = itertools.permutations(multiplied, 10)
-#         self.type = zp.DATASOURCE_TYPE.TRADE
-        
-#         self.events = (mock_data_unframe(id, 
         
         
 def mock_data_unframe(source_id, dt, type):

zipline/gens/tradegens.py

@@ -0,0 +1,40 @@
+import random
+from itertools import chain, repeat, cycle
+from datetime import datetime, timedelta
+
+from zipline.utils.factory import create_trade, create_trade
+from zipline.gens.utils import date_gen
+
+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"""
+
+    if rand:
+        return (random.uniform(0.0, 10.0) for i in xrange(n))
+    else:
+        return (float(i % 11) for i in xrange(1,n+1))
+
+def mock_volumes(n, rand = False):
+    """Does the same as mock_prices.  Different function name
+    for readability."""
+    return mock_prices(n, rand)
+
+def SpecificEquityTrades(n = 500, sids = [1, 2], event_list = None):
+    """Returns the first n events of event_list if specified. 
+    Otherwise generates a sensible stream of events."""
+    
+    if event_list:
+        return (event for event in event_list)
+    
+    else:
+        dates = date_gen(n = n)
+        prices = mock_prices(n)
+        volumes = mock_volumes(n)
+        sids = cycle(iter(sids))
+        
+        arg_gen = izip(sids, prices, volumes, dates)
+        
+        trades = (create_trade(*args) for args in arg_gen)
+        
+        return trades

zipline/gens/utils.py

@@ -48,7 +48,6 @@ def assert_datasource_protocol(event):
         assert isinstance(event.dt, datetime)
         assert event.dt.tzinfo == pytz.utc
     
-
 def assert_trade_protocol(event):
     """Assert that an event meets the protocol for datasource TRADE outputs."""
     assert_datasource_protocol(event)

zipline/utils/factory.py

@@ -89,6 +89,8 @@ def get_next_trading_dt(current, interval, trading_calendar):
 
     return next
 
+
+
 def create_trade_history(sid, prices, amounts, interval, trading_calendar):
     trades = []
     current = trading_calendar.first_open