diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py new file mode 100644 index 00000000..1b710212 --- /dev/null +++ b/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: + + + diff --git a/zipline/gens/merge.py b/zipline/gens/merge.py new file mode 100644 index 00000000..7c0a195c --- /dev/null +++ b/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 + + + + diff --git a/zipline/gens/test_feed.py b/zipline/gens/test_feed.py index 8e6a4ff0..5e632f9d 100644 --- a/zipline/gens/test_feed.py +++ b/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): diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py new file mode 100644 index 00000000..9be014e4 --- /dev/null +++ b/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 diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py index 8d168d8e..efe81a14 100644 --- a/zipline/gens/utils.py +++ b/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) diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index 6101c3b2..23c2eb3e 100644 --- a/zipline/utils/factory.py +++ b/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