mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-15 11:22:18 +08:00
Removes sort module.
Sort module is now unneeded with use of heapq.merge Also adapts test_perf_tracking so that it uses date_sorted_sources.
This commit is contained in:
Eddie Hebert
+16
-22
@@ -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,
|
||||
|
||||
@@ -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})
|
||||
@@ -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"
|
||||
@@ -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])
|
||||
|
||||
Reference in New Issue
Block a user