Get a subset of a generator

问题:

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?

回答1:

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.

回答2:

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))