Truncating a python generator

Question

I would like to know the Pythonic way to write a generator expression that takes the first n elements of an infinite generator g. Currently, I am doing this:

(x for x,_ in zip(g,range(n)))

Is there a more Pythonic way to do this?

Answer 1

islice

To wrap itertools.islice in a function would make a lot of sense, the code is much shorter than my alternative below:

from itertools import islice

def first_n(iterable, n):
    return islice(iterable, 0, n)

and usage:

>>> first_n(range(100), 10)
<itertools.islice object at 0xffec1dec>
>>> list(first_n(range(100), 10))
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

or directly:

>>> list(islice(range(100), 0, 10)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

alternative generator

here's an alternative:

def first_n(iterable, n):
    it = iter(iterable)
    for _ in range(n):
        yield next(it)

and usage:

>>> first_n(range(100), 10)
<generator object first_n at 0xffec73ec>
>>> list(first_n(range(100), 10))
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]