How to limit the size of a comprehension?
Same solution just without islice
:
filtered = (i for i in a if i == 1)
b = [filtered.next() for j in range(3)]
BTW, pay attention if the generator is empty or if it has less than 3 - you'll get StopIteration Exception.
To prevent that, you may want to use next() with default value. For example:
b = [next(filtered, None) for j in range(3)]
And if you don't want 'None' in your list:
b = [i for i in b if i is not None]
You could use itertools.count
to generate a counter and itertools.takewhile
to stop the iterating over a generator when the counter reaches the desired integer (3
in this case):
from itertools import count, takewhile
c = count()
b = list(takewhile(lambda x: next(c) < 3, (i for i in a if i == 1)))
Or a similar idea building a construct to raise StopIteration
to terminate the generator. That is the closest you'll get to your original idea of breaking the list comprehension, but I would not recommend it as best practice:
c = count()
b = list(i if next(c) < 3 else next(iter([])) for i in a if i == 1)
Examples:
>>> a = [1,2,1,4,1,1,1,1]
>>> c = count()
>>> list(takewhile(lambda x: next(c) < 3, (i for i in a if i == 1)))
[1, 1, 1]
>>> c = count()
>>> list(i if next(c) < 3 else next(iter([])) for i in a if i == 1)
[1, 1, 1]
@Martijn Pieters is completly right that itertools.islice
is the best way to solve this. However if you don't mind an additional (external) library you can use iteration_utilities
which wraps a lot of these itertools
and their applications (and some additional ones). It could make this a bit easier, at least if you like functional programming:
>>> from iteration_utilities import Iterable
>>> Iterable([1, 2, 1, 2, 1, 2]).filter((1).__eq__)[:2].as_list()
[1, 1]
>>> (Iterable([1, 2, 1, 2, 1, 2])
... .filter((1).__eq__) # like "if item == 1"
... [:2] # like "islice(iterable, 2)"
... .as_list()) # like "list(iterable)"
[1, 1]
The iteration_utilities.Iterable
class uses generators internally so it will only process as many items as neccessary until you call any of the as_*
(or get_*
) -methods.
Disclaimer: I'm the author of the iteration_utilities
library.
You can use a generator expression to do the filtering, then use islice()
to limit the number of iterations:
from itertools import islice
filtered = (i for i in a if i == 1)
b = list(islice(filtered, 3))
This ensures you don't do more work than you have to to produce those 3 elements.
Note that there is no point anymore in using a list comprehension here; a list comprehension can't be broken out of, you are locked into iterating to the end.