diff --git a/zipline/gens/feed.py b/zipline/gens/feed.py index cdc2f7b7..0d0b0bec 100644 --- a/zipline/gens/feed.py +++ b/zipline/gens/feed.py @@ -25,7 +25,6 @@ def FeedGen(stream_in, source_ids): message and yield it. """ - assert isinstance(stream_in, types.GeneratorType) assert isinstance(source_ids, list) # Set up an internal queue for each expected source. @@ -109,7 +108,7 @@ def pop_oldest(sources): oldest_event = older(oldest_event, current_event) # Pop the oldest event we found from its queue and return it. - return sources[oldest_event.source_id].pop() + return sources[oldest_event.source_id].popleft() # Return the event with the older timestamp. Break ties by source_id. def older(oldest, current): diff --git a/zipline/gens/test_feed.py b/zipline/gens/test_feed.py index 1605b6aa..8e6a4ff0 100644 --- a/zipline/gens/test_feed.py +++ b/zipline/gens/test_feed.py @@ -6,7 +6,7 @@ import pytz from unittest2 import TestCase from pymongo import Connection, ASCENDING -from itertools import izip, izip_longest, permutations, cycle +from itertools import izip, izip_longest, permutations, cycle, chain from datetime import datetime, timedelta from collections import deque @@ -14,7 +14,7 @@ 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 + assert_trade_protocol, date_gen, alternate import zipline.protocol as zp @@ -96,7 +96,8 @@ class FeedGenTestCase(TestCase): Assert that FeedGen's output agrees with expected. """ feed_gen = FeedGen(events, source_ids) - assert list(feed_gen) == expected + l = list(feed_gen) + assert l == expected def test_single_source(self): @@ -119,38 +120,41 @@ class FeedGenTestCase(TestCase): self.run_FeedGen(event_gen, expected, source_ids) - def test_multi_source_interleaved(self): + def test_multi_source(self): source_ids = ['a', 'b'] type = zp.DATASOURCE_TYPE.TRADE - # Set up source 'a'. Outputs 3 events with 2 minute deltas. + # 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 = 3)) + dates_a = list(date_gen(delta = delta_a, n = 20)) dates_a.append("DONE") events_a_args = zip(cycle(['a']), iter(dates_a), cycle([type])) events_a = [mock_data_unframe(*args) for args in events_a_args] - event_gen_a = (e for e in events_a) - - # Set up source 'b'. Outputs 4 events with 1 minute deltas. + + # 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 = 4)) + dates_b = list(date_gen(delta = delta_b, n = 10)) dates_b.append("DONE") events_b_args = zip(cycle(['b']), iter(dates_b), cycle([type])) events_b = [mock_data_unframe(*args) for args in events_b_args] - event_gen_b = (e for e in events_b) - # The expected output is all non-DONE events in both a and b, # sorted first by dt and then by source_id. non_dones = events_a[:-1] + events_b[:-1] expected = sorted(non_dones, compare_by_dt_source_id) - - import nose.tools; nose.tools.set_trace() - self.run_FeedGen(event_gen, expected, source_ids) - + # Alternating between a and b. + interleaved = alternate(iter(events_a), iter(events_b)) + self.run_FeedGen(interleaved, expected, source_ids) + + # All of a, then all of b. + + sequential = chain(iter(events_a), iter(events_b)) + self.run_FeedGen(sequential, expected, source_ids) + + # def test_FeedGen_consistency(self): # source_ids = ['a', 'b'] diff --git a/zipline/gens/test_mongods.py b/zipline/gens/test_mongods.py index 7260fd7d..ea19d90e 100644 --- a/zipline/gens/test_mongods.py +++ b/zipline/gens/test_mongods.py @@ -15,15 +15,6 @@ from zipline.gens.utils import stringify_args, assert_datasource_protocol,\ import zipline.protocol as zp -def mock_raw_event(sid, dt): - event = { - 'sid' : sid, - 'dt' : dt, - 'price' : 1.0, - 'volume' : 1 - } - return event - mongo_conn_args = { 'mongodb_host': 'localhost', 'mongodb_port': 27017, diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py index fd1c8464..1d9a6fd9 100644 --- a/zipline/gens/utils.py +++ b/zipline/gens/utils.py @@ -3,7 +3,7 @@ import numbers from hashlib import md5 from datetime import datetime, timedelta - +from itertools import izip_longest from zipline import ndict from zipline.protocol import DATASOURCE_TYPE @@ -19,6 +19,13 @@ def mock_raw_event(sid, dt): 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): + for e1, e2 in izip_longest(g1, g2): + if e1 != None: + yield e1 + if e2 != None: + yield e2 + def stringify_args(*args, **kwargs): """Define a unique string for any set of representable args.""" arg_string = '_'.join([str(arg) for arg in args])