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?

Asked By: Davis Yoshida

||

Answers:

itertools.islice(g,0,10)

should do it?

you might need list(itertools.islice(g,0,10)) (since it returns an iterator)

Answered By: Joran Beasley

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]

From the doc, itertools.islice has signature:

itertools.islice(iterable, stop)
itertools.islice(iterable, start, stop[, step])

So you could use:

g2 = itertools.islice(g, 10)


list(g2) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
list(g2) == []  # No elements after the iterable is consumed once

(no need to wrap itertools.islice, it can be used directly)

Answered By: Conchylicultor
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.