Generator that takes any number of sorted iterables and yields all values of all iterables in an ascending order

merge_iterators.py

import heapq
import unittest


def merge(*args):
    iterators = [iter(i) for i in args]
    heads = []
    for idx, iterator in enumerate(iterators):
        try:
            value = next(iterator)
            heapq.heappush(heads, (value, idx))
        except StopIteration:
            continue

    while heads:
        value, idx = heapq.heappop(heads)
        if heads:
            next_value = heads[0][0]
        else:
            next_value = float("inf")
        # idx is an index of an iterator with a min value
        # yield from that iterator
        while value <= next_value:
            yield value
            try:
                value = next(iterators[idx])
                if value > next_value:
                    heapq.heappush(heads, (value, idx))
            except StopIteration:
                break
                
class MergeTestCase(unittest.TestCase):
    def test_merge(self):
        cases = [
            (
                [
                    [1, 3, 5],
                    [2, 4, 6]
                ],
                [1, 2, 3, 4, 5, 6],
            ),
            (
                [],
                [],
            ),
            (
                [[], [], []],
                [],
            ),
            (
                [
                    [],
                    [1],
                    [0],
                    [],
                ],
                [0, 1],
            ),
            (
                [
                    [1, 3, 5],
                    [2, 2, 2],
                    [0, 0, 0],
                ],
                [0, 0, 0, 1, 2, 2, 2, 3, 5],
            ),
        ]

        for iterables, expected in cases:
            self.assertEqual(expected, list(merge(*iterables)))

    def test_big(self):
        result = []
        for i in range(1_000_000):
            result.append(i)
            result.append(i)

        self.assertEqual(result, list(merge(
            range(1_000_000),
            range(1_000_000),
        )))
        
unittest.main(verbosity=2)