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

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

Answer 2

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 3

Instead of a "selection function", I would use a "sorting function" which tells which element should go first.

The program starts by creating a list of 2-tuples: ( iterator, current value ). Since it is possible for one iterator to be empty, this has to be done with a try..catch (i.e. it cannot be in compact form).

Second, we iterate as long as there is at least one iterator. The sorting function put the element that has to go out in the first place. This element is "yielded". After that, the iterator is called to get the next element. If there is no more elements, the iterator is removed from the list.

That gives the following code

def iselect( list_of_iterators, sort_function ):
  work_list = []
  for i in list_of_iterators:
    try:
      new_item = ( i, next(i) )  # iterator and its first element
      work_list.append( new_item )
    except StopIteration:
      pass                      # this iterator is empty, skip it
  #
  while len(work_list) > 0:
    # this selects which element should go first
    work_list.sort( lambda e1,e2: sort_function(e1[1],e2[1]) )
    yield work_list[0][1]
    # update the first element of the list
    try:
      i, e = work_list[0]
      e = next(i)
      work_list[0] = ( i, e )
    except StopIteration:
      work_list = work_list[1:]

to test this program (including an iterator that yields nothing), I used

def iter_vowels():
  for v in 'aeiouy':
    yield v

def iter_consonnants():
  for c in 'bcdfghjklmnpqrstvwxz':
    yield c

def no_letters():
  if 1==2:     # not elegant, but.. 
    yield "?"  # .."yield" must appear to make this a generator

def test():
  i1 = iter_vowels()
  i2 = iter_consonnants()
  i3 = no_letters()
  sf = lambda x,y: cmp(x,y)
  for r in iselect( (i1,i2,i3), sf ):
    print (r)

test()

Answer 4

You could send it back into generator:

def iselect(i1, i2, f):
    while True:
        try:
            e1, e2 = next(i1), next(i2)
            if f(e1, e2):
                yield e1
                i2.send(e2)
            else:
                yield e2
                i1.send(e1)
        except StopIteration:
            return