diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index ba0538b6..5ee45de6 100644 --- a/zipline/gens/composites.py +++ b/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) diff --git a/zipline/gens/feed.py b/zipline/gens/feed.py index 284023a0..f029d0ee 100644 --- a/zipline/gens/feed.py +++ b/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()) ) diff --git a/zipline/gens/merge.py b/zipline/gens/merge.py index 7c0a195c..52adcb84 100644 --- a/zipline/gens/merge.py +++ b/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 diff --git a/zipline/gens/test_feed.py b/zipline/gens/test_feed.py index 7c486e33..b5007b91 100644 --- a/zipline/gens/test_feed.py +++ b/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): diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py index 35100f8e..fb0b3f48 100644 --- a/zipline/gens/tradegens.py +++ b/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) diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index 666fdd46..137b9d69 100644 --- a/zipline/gens/transform.py +++ b/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] diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py index 6150757b..c10cffdf 100644 --- a/zipline/gens/utils.py +++ b/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()) + diff --git a/zipline/gens/zmq_gens.py b/zipline/gens/zmq_gens.py new file mode 100644 index 00000000..524852a7 --- /dev/null +++ b/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)