diff --git a/Lib/heapq.py b/Lib/heapq.py index 00b429c..6cb22c3 100644 --- a/Lib/heapq.py +++ b/Lib/heapq.py @@ -357,29 +357,36 @@ def merge(*iterables): [0, 1, 2, 3, 4, 5, 5, 7, 8, 10, 15, 20, 25] ''' - _heappop, _heapreplace, _StopIteration = heappop, heapreplace, StopIteration + _heappop = heappop + _heapreplace = heapreplace + _StopIteration = StopIteration + _len = len h = [] h_append = h.append for itnum, it in enumerate(map(iter, iterables)): try: - next = it.__next__ - h_append([next(), itnum, next]) + itnext = it.__next__ + h_append([itnext(), itnum, itnext]) except _StopIteration: - pass - heapify(h) + pass # no items yielded at all - while 1: + heapify(h) + while _len(h) > 1: try: - while 1: - v, itnum, next = s = h[0] # raises IndexError when h is empty - yield v - s[0] = next() # raises StopIteration when exhausted - _heapreplace(h, s) # restore heap condition + while True: + value, itnum, itnext = s = h[0] + yield value + s[0] = itnext() # raises StopIteration when exhausted + _heapreplace(h, s) # restore heap condition except _StopIteration: - _heappop(h) # remove empty iterator - except IndexError: - return + _heappop(h) # remove empty iterator + + if h: + # fast case when only a single iterator remains + value, itnum, itnext = h[0] + yield value + yield from itnext.__self__ # Extend the implementations of nsmallest and nlargest to use a key= argument _nsmallest = nsmallest