Iterator selector in Python

Question

Is there a standard pythonic way of selecting a value from a list of provided iterators without advancing those that were not selected?

Something in the vein of this for two iterators (don't judge this too hard: it was quickly thrown together just to illustrate the idea):

def iselect(i1, i2, f):
    e1_read = False
    e2_read = False

    while True:
        try:
            if not e1_read:
                e1 = next(i1)
                e1_read = True

            if not e2_read:
                e2 = next(i2)
                e2_read = True

            if f(e1, e2):
                yield e1
                e1_read = False
            else:
                yield e2
                e2_read = False
        except StopIteration:
            return

Note that if one uses something like this instead:

[e1 if f(e1, e2) else e2 for (e1, e2) in zip(i1, i2)]

then the non-selected iterator advances every time, which is not what I want.

Answer 1

The more-itertools package has a peekable wrapper for iterators. It would seem like this should allow for a very clean solution if I understand your question correctly. You need to peek at the current values of a set of iterators and only modify the chosen iterator by calling next() on it.

from more_itertools import peekable

# the implementation of iselect can be very clean if
# the iterators are peekable
def iselect(peekable_iters, selector):
    """
    Parameters
    ----------
    peekable_iters: list of peekables
       This is the list of iterators which have been wrapped using
       more-itertools peekable interface.
    selector: function
       A function that takes a list of values as input, and returns
       the index of the selected value.
    """
    while True:
        peeked_vals = [it.peek(None) for it in peekable_iters]
        selected_idx = selector(peeked_vals)  # raises StopIteration
        yield next(peekable_iters[selected_idx])

Test this code:

# sample input iterators for testing
# assume python 3.x so range function returns iterable
iters = [range(i,5) for i in range(4)]

# the following could be encapsulated...
peekables = [peekable(it) for it in iters]

# sample selection function, returns index of minimum
# value among those being compared, or StopIteration if
# one of the lists contains None
def selector_func(vals_list):
    if None in vals_list:
        raise StopIteration
    else:
        return vals_list.index(min(vals_list))

for val in iselect(peekables, selector_func):
    print(val)    

Output:

0
1
1
2
2
2
3
3
3
3
4

Answer 2

You could use itertools.chain to prepend the last item back onto the iterator:

import itertools as IT
iterator = IT.chain([item], iterator)

And with many iterators:

items = map(next, iterators)
idx = f(*items)
iterators = [IT.chain([item], iterator) if i != idx else iterator
             for i, (item, iterator) in enumerate(zip(items, iterators))]

---

For example,

import itertools as IT

def iselect(f, *iterators):
    iterators = map(iter, iterators)
    while True:
        try:
            items = map(next, iterators)
        except StopIteration:
            return
        idx = f(*items)
        iterators = [IT.chain([item], iterator) if i != idx else iterator
                     for i, (item, iterator) in enumerate(zip(items, iterators))]
        yield items[idx]

def foo(*args):
    return sorted(range(len(args)), key=args.__getitem__)[0]

i1 = range(4)
i2 = range(4)
i3 = range(4)
for item in iselect(foo, i1, i2, i3):
    print(item)

yields

0
0
0
1
1
1
2
2
2
3