Get a subset of a generator
Question:
I have a generator function and want to get the first ten items from it; my first attempt was:
my_generator()[:10]
This doesn’t work because generators aren’t subscriptable, as the error tells me. Right now I have worked around that with:
list(my_generator())[:10]
This works since it converts the generator to a list; however, it’s inefficient and defeats the point of having a generator. Is there some built-in, Pythonic equivalent of [:10]
for generators?
Answers:
import itertools
itertools.islice(mygenerator(), 10)
itertools
has a number of utilities for working with iterators. islice
takes start, stop, and step arguments to slice an iterator just as you would slice a list.
to clarify the above comments:
from itertools import islice
def fib_gen():
a, b = 1, 1
while True:
yield a
a, b = b, a + b
assert [1, 1, 2, 3, 5] == list(islice(fib_gen(), 5))
I have a generator function and want to get the first ten items from it; my first attempt was:
my_generator()[:10]
This doesn’t work because generators aren’t subscriptable, as the error tells me. Right now I have worked around that with:
list(my_generator())[:10]
This works since it converts the generator to a list; however, it’s inefficient and defeats the point of having a generator. Is there some built-in, Pythonic equivalent of [:10]
for generators?
import itertools
itertools.islice(mygenerator(), 10)
itertools
has a number of utilities for working with iterators. islice
takes start, stop, and step arguments to slice an iterator just as you would slice a list.
to clarify the above comments:
from itertools import islice
def fib_gen():
a, b = 1, 1
while True:
yield a
a, b = b, a + b
assert [1, 1, 2, 3, 5] == list(islice(fib_gen(), 5))