How to return the first n items of an iterator that pass a test without running the test on the remainder

Question

I'm looking for something like:

foo = [ p for p in somelist if xxx(p)][0:10]

This works, but executes xxx on all of somelist, and xxx is expensive.

I could write a subfunction such as

def blah (list, len): 
    res=[]; i=0; 
    while i<len: 
        if xxx(list[i]): 
            res.append(i) 
   return res

but this seems very unpythonic.

An example that proves the first result does the wrong thing would be:

foo = [ p for p in (1,2,3,4,5,0,6) if 100/p ][0:3]

which "should" return [1,2,3] but in fact fails on the divide by zero.

I've tried various tools from itertools but can't find a combination that stops execution of the iterator after the size is reached.

Answer 1

Try itertools.islice:

from itertools import islice
foo = list(islice((p for p in (1,2,3,4,5,0,6) if 100/p), 4))

Notice the lack of brackets: this is a generator comprehension

Answer 2

That is impossible to do with list comprehensions but it is possible with generator comprehensions.

What is the difference between the two? List comprehensions will iterate over the iterable all at once and return a list back to you according to your operations on items on the iterable. The key here is all at once.

Your example uses list comprehensions and here's what happens: First the list comprehension is evaluated. In your case it will fail, but even if it didn't fail, it would have iterated over everything in the iterable somelist and would have returned a resulting list back. Then this list is sliced and a new resulting list is returned.

Generator comprehensions and generators in general have different way of working. They are basically code blocks that are suspended until you request more data from them, which is what you really want to do.

In effect, you create a generator as follows:

g = (p for p in (1,2,3,4,5,0,6) if 100/p )

You now have a generator that will generate values for you when you request it to do so, according to the "rule" you gave it.

As soon as you have a generator at hand, there are several ways you could get n items from it.

You could write a simple for loop as follows:

results = []
for x in range(3): # n = 3 in your case
    results.append(next(g))

Of course, that isn't Pythonic. Assuming you want a list back, you can now use a list comprehension:

results = [next(g) for x in range(3)]

This is the "manual" way of doing it. You could also use the islice function from the itertools module (Documentation here):

import itertools
results = list(itertools.islice(g, 4))

And there it is. Generators are pretty useful. A code block that executes when requested to and remember its state is truly invaluable.

Answer 3

By using a combination of the builtin filter and islice you can achieve what you want. eg.

length = 3
items_iter = islice(filter(None, [1, 2, 0, 4, 5]), length) # returns a lazy iterator
items = list(items_iter)
assert items == [1, 2, 5]

However, you may wish to write your own generator. Generator functions will, rather than returning a single result, will yield successive items. Sometimes it might yield no items, and sometimes they can yield an infinite number of items. Here is an example for your case:

def take_first_n_that_pass(test, length, iterator):
    i = 0
    for item in iterator:
        if i >= length:
            return # stops generator
        if test(item):
            yield item
            i += 1
    # end of iterator -- generator stops here as well

items_iter = take_first_n_that_pass(bool, 3, [0, 1, 2, 0, 4, 5])
items = list(items_iter)
assert items == [1, 2, 4]