mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-10 20:34:20 +08:00
Merge pull request #53 from quantopian/performance-improvements-sorting-and-done-message
Performance improvements sorting and done message
This commit is contained in:
Richard Frank
+13
-23
@@ -27,8 +27,8 @@ 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.gens.composites import date_sorted_sources
|
||||
|
||||
from zipline.finance.trading import TradingEnvironment
|
||||
|
||||
@@ -642,8 +642,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 +657,14 @@ 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]
|
||||
|
||||
perf_messages = \
|
||||
[msg for date, snapshot in
|
||||
perf_tracker.transform(
|
||||
@@ -685,12 +672,15 @@ class TestPerformanceTracker(unittest.TestCase):
|
||||
for event in snapshot
|
||||
for msg in event.perf_messages]
|
||||
|
||||
end_perf_messages, risk_message = perf_tracker.handle_simulation_end()
|
||||
|
||||
perf_messages.extend(end_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,
|
||||
@@ -702,7 +692,7 @@ class TestPerformanceTracker(unittest.TestCase):
|
||||
def event_with_txn(self, event, no_txn_dt):
|
||||
#create a transaction for all but
|
||||
#first trade in each sid, to simulate None transaction
|
||||
if event.dt != no_txn_dt and event.dt != 'DONE':
|
||||
if event.dt != no_txn_dt:
|
||||
txn = ndict({
|
||||
'sid': event.sid,
|
||||
'amount': -25,
|
||||
|
||||
@@ -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})
|
||||
@@ -219,12 +219,8 @@ class PerformanceTracker(object):
|
||||
new_snapshot = []
|
||||
|
||||
for event in snapshot:
|
||||
if date != "DONE":
|
||||
event.perf_messages = self.process_event(event)
|
||||
event.portfolio = self.get_portfolio()
|
||||
else:
|
||||
event.perf_messages, event.risk_message = \
|
||||
self.handle_simulation_end()
|
||||
event.perf_messages = self.process_event(event)
|
||||
event.portfolio = self.get_portfolio()
|
||||
|
||||
del event['TRANSACTION']
|
||||
new_snapshot.append(event)
|
||||
|
||||
+10
-26
@@ -13,10 +13,12 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from itertools import chain
|
||||
import heapq
|
||||
|
||||
from zipline.gens.utils import roundrobin, done_message
|
||||
from zipline.gens.sort import date_sort
|
||||
|
||||
def _decorate_source(source):
|
||||
for message in source:
|
||||
yield ((message.dt, message.source_id), message)
|
||||
|
||||
|
||||
def date_sorted_sources(*sources):
|
||||
@@ -24,23 +26,11 @@ def date_sorted_sources(*sources):
|
||||
Takes an iterable of sources, generating namestrings and
|
||||
piping their output into date_sort.
|
||||
"""
|
||||
sorted_stream = heapq.merge(*(_decorate_source(s) for s in sources))
|
||||
|
||||
for source in sources:
|
||||
assert iter(source), "Source %s not iterable" % source
|
||||
assert hasattr(source, 'get_hash'), "No get_hash"
|
||||
|
||||
# Get name hashes to pass to date_sort.
|
||||
names = [source.get_hash() for source in sources]
|
||||
|
||||
# Convert the list of generators into a flat stream by pulling
|
||||
# one element at a time from each.
|
||||
stream_in = roundrobin(sources, names)
|
||||
|
||||
# Guarantee the flat stream will be sorted by date, using
|
||||
# source_id as tie-breaker, which is fully deterministic (given
|
||||
# deterministic string representation for all args/kwargs)
|
||||
|
||||
return date_sort(stream_in, names)
|
||||
# Strip out key decoration
|
||||
for _, message in sorted_stream:
|
||||
yield message
|
||||
|
||||
|
||||
def sequential_transforms(stream_in, *transforms):
|
||||
@@ -61,8 +51,7 @@ def sequential_transforms(stream_in, *transforms):
|
||||
transforms,
|
||||
stream_in)
|
||||
|
||||
dt_aliased = alias_dt(stream_out)
|
||||
return add_done(dt_aliased)
|
||||
return alias_dt(stream_out)
|
||||
|
||||
|
||||
def alias_dt(stream_in):
|
||||
@@ -72,8 +61,3 @@ def alias_dt(stream_in):
|
||||
for message in stream_in:
|
||||
message['datetime'] = message['dt']
|
||||
yield message
|
||||
|
||||
|
||||
# Add a done message to a stream.
|
||||
def add_done(stream_in):
|
||||
return chain(stream_in, [done_message('Composite')])
|
||||
|
||||
@@ -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"
|
||||
@@ -71,6 +71,7 @@ class TradeSimulationClient(object):
|
||||
self.algo_start = self.environment.first_open
|
||||
self.algo_sim = AlgorithmSimulator(
|
||||
self.ordering_client,
|
||||
self.perf_tracker,
|
||||
self.algo,
|
||||
self.algo_start
|
||||
)
|
||||
@@ -116,6 +117,7 @@ class AlgorithmSimulator(object):
|
||||
|
||||
def __init__(self,
|
||||
order_book,
|
||||
perf_tracker,
|
||||
algo,
|
||||
algo_start):
|
||||
|
||||
@@ -126,6 +128,7 @@ class AlgorithmSimulator(object):
|
||||
# We extract the order book from the txn client so that
|
||||
# the algo can place new orders.
|
||||
self.order_book = order_book
|
||||
self.perf_tracker = perf_tracker
|
||||
|
||||
self.algo = algo
|
||||
self.algo_start = algo_start.replace(hour=0, minute=0,
|
||||
@@ -203,18 +206,10 @@ class AlgorithmSimulator(object):
|
||||
if self.simulation_dt is None:
|
||||
self.simulation_dt = date
|
||||
|
||||
# Done message has the risk report, so we yield before exiting.
|
||||
if date == 'DONE':
|
||||
for event in snapshot:
|
||||
for perf_message in event.perf_messages:
|
||||
yield perf_message
|
||||
yield event.risk_message
|
||||
raise StopIteration
|
||||
|
||||
# We're still in the warmup period. Use the event to
|
||||
# update our universe, but don't yield any perf messages,
|
||||
# and don't send a snapshot to handle_data.
|
||||
elif date < self.algo_start:
|
||||
if date < self.algo_start:
|
||||
for event in snapshot:
|
||||
del event['perf_messages']
|
||||
self.update_universe(event)
|
||||
@@ -233,6 +228,14 @@ class AlgorithmSimulator(object):
|
||||
# to the user's algo.
|
||||
self.simulate_snapshot(date)
|
||||
|
||||
perf_messages, risk_message = \
|
||||
self.perf_tracker.handle_simulation_end()
|
||||
|
||||
for message in perf_messages:
|
||||
yield message
|
||||
|
||||
yield risk_message
|
||||
|
||||
def update_universe(self, event):
|
||||
"""
|
||||
Update the universe with new event information.
|
||||
|
||||
@@ -17,11 +17,9 @@
|
||||
import pytz
|
||||
import numbers
|
||||
|
||||
from collections import OrderedDict
|
||||
from hashlib import md5
|
||||
from datetime import datetime
|
||||
from itertools import izip_longest
|
||||
from zipline import ndict
|
||||
from zipline.protocol import (
|
||||
DATASOURCE_TYPE,
|
||||
Event
|
||||
@@ -38,18 +36,6 @@ def mock_raw_event(sid, dt):
|
||||
return event
|
||||
|
||||
|
||||
def mock_done(id):
|
||||
return ndict({
|
||||
'dt': "DONE",
|
||||
"source_id": id,
|
||||
'tnfm_id': id,
|
||||
'tnfm_value': None,
|
||||
'type': DATASOURCE_TYPE.DONE
|
||||
})
|
||||
|
||||
done_message = mock_done
|
||||
|
||||
|
||||
def alternate(g1, g2):
|
||||
"""Specialized version of roundrobin for just 2 generators."""
|
||||
for e1, e2 in izip_longest(g1, g2):
|
||||
@@ -59,27 +45,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