Sortable itertools.chain with key.

sorted_chain.py

"""
@description: Sortable chain with key, just like sorted.
@author: Laisky
@copyright: 2015/Laisky
@license: MIT/Apache
"""

from collections import namedtuple


def sorted_chain(iterators, key=None):
    _iter_map = dict(zip(iterators, [None] * len(iterators)))
    _middle_iter = namedtuple('middle_iter', ('iterator', 'value'))
    
    while _iter_map:
        _tmp = None
        _current = None
        _exhausted = []
        for _iter, val in _iter_map.items():
            if val is None:
                try:
                    val = next(_iter)
                except StopIteration:
                    _exhausted.append(_iter)
                    continue
                else:
                    _iter_map[_iter] = val
                
            if _current is None:
                _current = _middle_iter(_iter, val)
            elif _tmp is None:
                _tmp = _middle_iter(_iter, val)
            
            if _current is not None and _tmp is not None:
                if key:
                    _is_change = key(_tmp.value) < key(_current.value)
                else:
                    _is_change = _tmp.value < _current.value
                
                if _is_change:                    
                    _current = _tmp
                
                _tmp = None
                   
        if _current:
            _iter_map[_current.iterator] = None
            yield _current.value
        
        for k in _exhausted:
            del _iter_map[k]
    
    
if __name__ == '__main__':
    it1 = iter([1, 3, 5, 7])
    it2 = iter([2, 4, 6, 8])
    it3 = iter([3,5,7,9])
    print(list(sorted_chain([it1, it2, it3])))
    # [1, 2, 3, 3, 4, 5, 5, 6, 7, 7, 8, 9]