mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-19 11:22:06 +08:00
Eddie Hebert
403a8e9064
Since this was invoked as an function, the invocation of the method was not dropped by -O. Also, moving towards only having postconditional asserts.
117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
#
|
|
# 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"
|