Random python functions for working with iterators.

iter_utils.py

"""
A number of functions for working with iterators.
"""

def all(iterable):
    for element in iterable:
        if not element:
            return False
    return True

def super_min(elems, key=None):
    """
    Takes a list of elements and an optional key and finds the smallest
    element. Returns the index, item tuple of the smallest elemnt or
    None, None if elems is empty.

    Ignores all elements that are None, and doesn't even apply key to them.
    """
    # get the values to compare
    if key is not None:
        keys = []
        for x in elems:
            if x is not None:
                keys.append(key(x))
            else:
                keys.append(x)
    else:
        keys = elems

    # find the smallest element
    m = None
    ind = None
    for i, item in enumerate(keys):
        if item is not None and (m is None or item < m):
            m = item
            ind = i

    # return the index and smallest element
    if m is None:
        return None, None
    else:
        return ind, elems[ind]

def iter_sorted_combine(iters, key=None):
    """
    Iterate in order over a number of other sorted iterators. Essentually
    implement the merge step in the merge sort algorithm using.

    Doesn't support iterating over elements that are None.

    `iters`: One or more iterators.
    `key`: An optional function to get the key to compare.
    """
    # prime all the iters
    elems = []
    for i in iters:
        try:
            elems.append(i.next())
        except StopIteration:
            elems.append(None)

    # until all the elements are None
    while not all(map(lambda x: x is None, elems)):
        # get the min value
        i, v = super_min(elems, key=key)

        # advance the iterator of the item
        try:
            elems[i] = iters[i].next()
        except StopIteration:
            elems[i] = None

        # yield the min item
        yield v

# Unit tests
import unittest
class TestSuperMin(unittest.TestCase):
    def test_empty(self):
        """Test the super_min of an empty list"""
        result = super_min([], key=lambda x: None)
        self.assertEqual(result, (None, None))

    def test_one(self):
        """Test a set of one element"""
        result = super_min([42])
        self.assertEqual(result, (0, 42))

    def test_many(self):
        """Test a bunch of elements"""
        result = super_min(['a', 'g', 'funk'])
        self.assertEqual(result, (0, 'a'))

    def test_with_key(self):
        """Test a bunch of elements with keys"""
        result = super_min([34, 234, 22, 34, 56, 93, 12.22222, 13], key=int)
        self.assertEqual(result, (6, 12.22222))

class TestIterSortedCombine(unittest.TestCase):
    def test_empty(self):
        """Test an empty iterator"""
        result = list(iter_sorted_combine([iter([])]))
        self.assertEqual(result, [])

    def test_one_element(self):
        """Test a set of one element"""
        result = list(iter_sorted_combine([iter([2]), iter([])]))
        self.assertEqual(result, [2])

    def test_many(self):
        """Test a bunch of elements, and three iterators"""
        result = list(iter_sorted_combine([
            iter([1, 3, 5, 8, 10]),
            iter([2, 9, 15]),
            iter([-3, 500])
        ]))
        self.assertEqual(result, [-3, 1, 2, 3, 5, 8, 9, 10, 15, 500])

    def test_with_key(self):
        """Test a bunch of elements with keys"""
        result = list(iter_sorted_combine([
            iter([(45, 1), (1, 2)]),
        ], key=lambda x: x[1]))
        self.assertEqual(result, [(45, 1), (1, 2)])

if __name__ == '__main__':
    # run all the tests independantly
    def test(typ):
        suite = unittest.TestLoader().loadTestsFromTestCase(typ)
        unittest.TextTestRunner(verbosity=2).run(suite)

    test(TestSuperMin)
    test(TestIterSortedCombine)