CODEWITHSUNDEEP
pythongenerator
Davis Yoshida
Davis Yoshida

Reputation: 1785

Truncating a python generator

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?

Upvotes: 4

Views: 1282

Answers (3)

Conchylicultor
Conchylicultor

Reputation: 5729

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)

Upvotes: 1

Aaron Hall
Aaron Hall

Reputation: 395125

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]

Upvotes: 9

Joran Beasley
Joran Beasley

Reputation: 113988

itertools.islice(g,0,10)

should do it?

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

Upvotes: 3

Related Questions