diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py index 76698f28..7cab2518 100644 --- a/tests/test_perf_tracking.py +++ b/tests/test_perf_tracking.py @@ -27,8 +27,9 @@ from operator import attrgetter import zipline.utils.factory as factory import zipline.finance.performance as perf from zipline.utils.protocol_utils import ndict -from zipline.gens.sort import date_sort -from zipline.protocol import DATASOURCE_TYPE +from zipline.protocol import Event + +from zipline.gens.composites import date_sorted_sources from zipline.finance.trading import TradingEnvironment @@ -642,8 +643,6 @@ class TestPerformanceTracker(unittest.TestCase): del trade_history[-days_to_delete.end:] del trade_history2[-days_to_delete.end:] - trade_history.extend(trade_history2) - trading_environment.first_open = \ trading_environment.calculate_first_open() trading_environment.last_close = \ @@ -659,25 +658,21 @@ class TestPerformanceTracker(unittest.TestCase): trading_environment ) - # date_sort requires 'DONE' messages from each source - events = itertools.chain(trade_history, - [ndict({ - 'source_id': 'factory1', - 'dt': 'DONE', - 'type': DATASOURCE_TYPE.TRADE - }), - ndict({ - 'source_id': 'factory2', - 'dt': 'DONE', - 'type': DATASOURCE_TYPE.TRADE - })]) - events = date_sort(events, ('factory1', 'factory2')) - events = itertools.chain(events, - [ndict({'dt': 'DONE'})]) + events = date_sorted_sources(trade_history, trade_history2) events = [self.event_with_txn(event, trade_history[0].dt) for event in events] + # Extract events with transactions to use for verification. + events_with_txns = [event for event in events if event.TRANSACTION] + + done_message = Event({ + 'dt': 'DONE', + 'TRANSACTION': None + }) + + events = itertools.chain(events, [done_message]) + perf_messages = \ [msg for date, snapshot in perf_tracker.transform( @@ -686,11 +681,10 @@ class TestPerformanceTracker(unittest.TestCase): for msg in event.perf_messages] #we skip two trades, to test case of None transaction - txn_count = len(trade_history) - 2 - self.assertEqual(perf_tracker.txn_count, txn_count) + self.assertEqual(perf_tracker.txn_count, len(events_with_txns)) cumulative_pos = perf_tracker.cumulative_performance.positions[sid] - expected_size = txn_count / 2 * -25 + expected_size = len(events_with_txns) / 2 * -25 self.assertEqual(cumulative_pos.amount, expected_size) self.assertEqual(perf_tracker.last_close, diff --git a/tests/test_sorting.py b/tests/test_sorting.py deleted file mode 100644 index a1034bf1..00000000 --- a/tests/test_sorting.py +++ /dev/null @@ -1,264 +0,0 @@ -# -# Copyright 2012 Quantopian, Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import pytz - -from unittest import TestCase -from itertools import chain, izip_longest -from datetime import datetime, timedelta -from collections import deque - -from zipline import ndict -from zipline.gens.sort import ( - date_sort, - done, - queue_is_done -) -from zipline.gens.utils import alternate, done_message -from zipline.sources import SpecificEquityTrades -from zipline.gens.composites import date_sorted_sources - - -class HelperTestCase(TestCase): - - def setUp(self): - pass - - def tearDown(self): - pass - - def test_individual_queue_logic(self): - queue = deque() - # Empty queues are neither done nor ready. - assert not queue_is_done(queue) - - queue.append(to_dt('foo')) - assert not queue_is_done(queue) - - queue.appendleft(to_dt('DONE')) - - # Checking done when we have a message after done will trip an assert. - self.assertRaises(AssertionError, queue_is_done, queue) - - queue.pop() - assert queue_is_done(queue) - - def test_pop_logic(self): - sources = {} - ids = ['a', 'b', 'c'] - for id in ids: - sources[id] = deque() - - assert not done(sources) - - # All sources must have a message to be ready/done - sources['a'].append(to_dt("datetime")) - assert not done(sources) - sources['a'].pop() - - for id in ids: - sources[id].append(to_dt("datetime")) - - assert not done(sources) - - for id in ids: - sources[id].appendleft(to_dt("DONE")) - - # ["DONE", message] will trip an assert in queue_is_done. - self.assertRaises(AssertionError, done, sources) - - for id in ids: - sources[id].pop() - - assert done(sources) - - -class DateSortTestCase(TestCase): - - def setUp(self): - pass - - def tearDown(self): - pass - - def run_date_sort(self, event_stream, expected, source_ids): - """ - Take a list of events, their source_ids, and an expected sorting. - Assert that date_sort's output agrees with expected. - """ - sort_out = date_sort(event_stream, source_ids) - for m1, m2 in izip_longest(sort_out, expected): - assert m1 == m2 - - def test_single_source(self): - - # Just using the built-in defaults. See - # zipline.sources.py - source = SpecificEquityTrades() - expected = list(source) - source.rewind() - # The raw source doesn't handle done messaging, so we need to - # append a done message for sort to work properly. - with_done = chain(source, [done_message(source.get_hash())]) - self.run_date_sort(with_done, expected, [source.get_hash()]) - - def test_multi_source(self): - - filter = [2, 3] - args_a = tuple() - kwargs_a = { - 'count': 100, - 'sids': [1, 2, 3], - 'start': datetime(2012, 1, 3, 15, tzinfo=pytz.utc), - 'delta': timedelta(minutes=6), - 'filter': filter - } - source_a = SpecificEquityTrades(*args_a, **kwargs_a) - - args_b = tuple() - kwargs_b = { - 'count': 100, - 'sids': [2, 3, 4], - 'start': datetime(2012, 1, 3, 15, tzinfo=pytz.utc), - 'delta': timedelta(minutes=5), - 'filter': filter - } - source_b = SpecificEquityTrades(*args_b, **kwargs_b) - - all_events = list(chain(source_a, source_b)) - - # The expected output is all events, sorted by dt with - # source_id as a tiebreaker. - expected = sorted(all_events, comp) - source_ids = [source_a.get_hash(), source_b.get_hash()] - - # Generating the events list consumes the sources. Rewind them - # for testing. - source_a.rewind() - source_b.rewind() - - # Append a done message to each source. - with_done_a = chain(source_a, [done_message(source_a.get_hash())]) - with_done_b = chain(source_b, [done_message(source_b.get_hash())]) - - interleaved = alternate(with_done_a, with_done_b) - - # Test sort with alternating messages from source_a and - # source_b. - self.run_date_sort(interleaved, expected, source_ids) - - source_a.rewind() - source_b.rewind() - with_done_a = chain(source_a, [done_message(source_a.get_hash())]) - with_done_b = chain(source_b, [done_message(source_b.get_hash())]) - - sequential = chain(with_done_a, with_done_b) - - # Test sort with all messages from a, followed by all messages - # from b. - - self.run_date_sort(sequential, expected, source_ids) - - def test_sort_composite(self): - - filter = [1, 2] - - #Set up source a. One hour between events. - args_a = tuple() - kwargs_a = { - 'count': 100, - 'sids': [1], - 'start': datetime(2012, 6, 6, 0), - 'delta': timedelta(hours=1), - 'filter': filter - } - source_a = SpecificEquityTrades(*args_a, **kwargs_a) - - #Set up source b. One day between events. - args_b = tuple() - kwargs_b = { - 'count': 50, - 'sids': [2], - 'start': datetime(2012, 6, 6, 0), - 'delta': timedelta(days=1), - 'filter': filter - } - source_b = SpecificEquityTrades(*args_b, **kwargs_b) - - #Set up source c. One minute between events. - args_c = tuple() - kwargs_c = { - 'count': 150, - 'sids': [1, 2], - 'start': datetime(2012, 6, 6, 0), - 'delta': timedelta(minutes=1), - 'filter': filter - } - source_c = SpecificEquityTrades(*args_c, **kwargs_c) - # Set up source d. This should produce no events because the - # internal sids don't match the filter. - args_d = tuple() - kwargs_d = { - 'count': 50, - 'sids': [3], - 'start': datetime(2012, 6, 6, 0), - 'delta': timedelta(minutes=1), - 'filter': filter - } - source_d = SpecificEquityTrades(*args_d, **kwargs_d) - sources = [source_a, source_b, source_c, source_d] - hashes = [source.get_hash() for source in sources] - - sort_out = date_sorted_sources(*sources) - - # Read all the values from sort and assert that they arrive in - # the correct sorting with the expected hash values. - to_list = list(sort_out) - copy = to_list[:] - - # We should have 300 events (100 from a, 150 from b, 50 from c) - assert len(to_list) == 300 - - for e in to_list: - # All events should match one of our expected source_ids. - assert e.source_id in hashes - # But none of them should match source_d. - assert e.source_id != source_d.get_hash() - - # The events should be sorted by dt, with source_id as tiebreaker. - expected = sorted(copy, comp) - - assert to_list == expected - - -def compare_by_dt_source_id(x, y): - if x.dt < y.dt: - return -1 - elif x.dt > y.dt: - return 1 - - elif x.source_id < y.source_id: - return -1 - elif x.source_id > y.source_id: - return 1 - else: - return 0 - -#Alias for ease of use -comp = compare_by_dt_source_id - - -def to_dt(msg): - return ndict({'dt': msg}) diff --git a/zipline/gens/sort.py b/zipline/gens/sort.py deleted file mode 100644 index 19e33cb5..00000000 --- a/zipline/gens/sort.py +++ /dev/null @@ -1,116 +0,0 @@ -# -# Copyright 2012 Quantopian, Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Sorting generator. -""" -import logbook - -from collections import deque - -log = logbook.Logger('Sorting') - - -def date_sort(stream_in, source_ids): - """ - A generator that takes a generator and a list of source_ids. We - maintain an internal queue for each id in source_ids. While we - have messages pending from all sources, we pull the earliest - message and yield it. - """ - assert isinstance(source_ids, (list, tuple)) - - # Set up an internal queue for each expected source. - sources = {} - for id in source_ids: - assert isinstance(id, basestring), "Bad source_id %s" % id - sources[id] = deque() - - # Process incoming streams. - for message in stream_in: - # Only allow messages from sources we expect. - assert message.source_id in sources, "Unexpected source: %s" % message - - sources[message.source_id].append(message) - - # Only pop messages when we have a pending message from - # all datasources. Stop if all sources have signalled done. - - while all(sources.values()) and not done(sources): - message = pop_oldest(sources) - 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 date_sort on exit: %s" % queue - assert queue[0].dt == "DONE", \ - "Bad last message in date_sort on exit: %s" % queue - - -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 date_sort: %s" % queue - return True - else: - return False - - -def pop_oldest(sources): - - oldest_event = None - - # Iterate over the dict, checking internal queues for the oldest - # pending event. - - for queue in sources.itervalues(): - current_event = queue[0] - # Skip queues that are done. - if current_event.dt == "DONE": - continue - # Any event is older than nothing. - elif oldest_event is None: - oldest_event = current_event - # Keep the older event. Break ties by source_id. This will - # trip an assert if we have duplicate sources. - else: - 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].popleft() - - -# Return the event with the older timestamp. Break ties by source_id. -def older(oldest, current): - # Try to compare by dt. - if oldest.dt < current.dt: - return oldest - elif oldest.dt > current.dt: - return current - # Break ties by source_id. - elif oldest.source_id < current.source_id: - return oldest - elif oldest.source_id > current.source_id: - return current - else: - assert False, "Duplicate event" diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py index c3b08cb4..2cdc2581 100644 --- a/zipline/gens/utils.py +++ b/zipline/gens/utils.py @@ -17,7 +17,6 @@ import pytz import numbers -from collections import OrderedDict from hashlib import md5 from datetime import datetime from itertools import izip_longest @@ -59,27 +58,6 @@ def alternate(g1, g2): yield e2 -def roundrobin(sources, namestrings): - """ - Takes N generators, pulling one element off each until all inputs - are empty. - """ - assert len(sources) == len(namestrings) - mapping = OrderedDict(zip(namestrings, sources)) - - # While our generators have not been exhausted, pull elements - while mapping.keys() != []: - for namestring, source in mapping.iteritems(): - try: - message = source.next() - # allow sources to yield None to avoid blocking. - if message: - yield message - except StopIteration: - yield done_message(namestring) - del mapping[namestring] - - def hash_args(*args, **kwargs): """Define a unique string for any set of representable args.""" arg_string = '_'.join([str(arg) for arg in args])