1.1 --- a/itermerge.py Wed Sep 09 01:18:04 2009 +0200
1.2 +++ b/itermerge.py Thu Sep 10 23:19:13 2009 +0200
1.3 @@ -1,54 +1,51 @@
1.4 #!/usr/bin/env python
1.5
1.6 -def itermerge(seq1, seq2):
1.7 +"An iterator merging class similar to heapq.merge in Python 2.6."
1.8
1.9 - "Merge 'seq1' and 'seq2' to produce an ordered, combined list of results."
1.10 +from bisect import insort_right
1.11
1.12 - results = []
1.13 +class itermerge:
1.14
1.15 - iter1 = iter(seq1)
1.16 - iter2 = iter(seq2)
1.17 + """
1.18 + Merge ordered sequences to produce an ordered, combined sequence of
1.19 + results.
1.20 + """
1.21
1.22 - t1 = None
1.23 - t2 = None
1.24 + def __init__(self, sequences):
1.25 + self.iters = []
1.26 +
1.27 + # Prepare the underlying iterators.
1.28 +
1.29 + for seq in sequences:
1.30 + it = iter(seq)
1.31 + next = it.next
1.32 + self._add_next(next)
1.33
1.34 - t1 = _itermerge_next(iter1)
1.35 - if t1 is None:
1.36 - _itermerge_fill(iter2, results)
1.37 - return results
1.38 + def sort(self):
1.39 + pass # The output should be sorted.
1.40 +
1.41 + def __iter__(self):
1.42 + return self
1.43 +
1.44 + def _add_next(self, next):
1.45
1.46 - while 1:
1.47 - if t1 is None:
1.48 - t1 = _itermerge_next(iter1)
1.49 - if t1 is None:
1.50 - results.append(t2)
1.51 - _itermerge_fill(iter2, results)
1.52 - return results
1.53 + """
1.54 + Store the current value for an iterator, alongside the means of
1.55 + getting the next value: the 'next' method.
1.56 + """
1.57 +
1.58 + try:
1.59 + insort_right(self.iters, (next(), next))
1.60 + except StopIteration:
1.61 + pass
1.62 +
1.63 + def next(self):
1.64 + if self.iters:
1.65 + value, next = self.iters[0]
1.66 + del self.iters[0]
1.67 + self._add_next(next)
1.68 + return value
1.69 else:
1.70 - t2 = _itermerge_next(iter2)
1.71 - if t2 is None:
1.72 - results.append(t1)
1.73 - _itermerge_fill(iter1, results)
1.74 - return results
1.75 -
1.76 - if t1 < t2:
1.77 - results.append(t1)
1.78 - t1 = None
1.79 - else:
1.80 - results.append(t2)
1.81 - t2 = None
1.82 -
1.83 -def _itermerge_next(iter):
1.84 - try:
1.85 - return iter.next()
1.86 - except StopIteration:
1.87 - return None
1.88 -
1.89 -def _itermerge_fill(iter, results):
1.90 - try:
1.91 - while 1:
1.92 - results.append(iter.next())
1.93 - except StopIteration:
1.94 - pass
1.95 + raise StopIteration
1.96
1.97 # vim: tabstop=4 expandtab shiftwidth=4