diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index ac24c7ce..1f68d58e 100644 --- a/zipline/gens/composites.py +++ b/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: - - + diff --git a/zipline/gens/feed.py b/zipline/gens/feed.py index e7498d5c..93e03371 100644 --- a/zipline/gens/feed.py +++ b/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. diff --git a/zipline/gens/test_feed.py b/zipline/gens/test_feed.py index 5e632f9d..9b77f9c8 100644 --- a/zipline/gens/test_feed.py +++ b/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 diff --git a/zipline/gens/test_mongods.py b/zipline/gens/test_mongods.py index ea19d90e..d9b8dbe5 100644 --- a/zipline/gens/test_mongods.py +++ b/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' diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py index 619721e4..8327aa63 100644 --- a/zipline/gens/tradegens.py +++ b/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]) diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py index 1060cd24..8966c638 100644 --- a/zipline/gens/utils.py +++ b/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]) diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index 23c2eb3e..db440891 100644 --- a/zipline/utils/factory.py +++ b/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,