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?

Asked By: Amandasaurus

||

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.

Answered By: Ned Batchelder

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))
Answered By: Dustin Getz
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.