Jean-Francois T.
Jean-Francois T.

Reputation: 12940

Stopping a Python iterator/generator after a given number of times

I have a generator with many elements, say

long_generator = (i**2 for i in range(10**1000))

I would like to extract the first n elements (without obviously parsing the generator until the end): what pythonic way could do this?

The function iter has a second parameter being a sentinel based on the returned value:

numbers = iter(lambda:next(long_generator), 81)  # but this assumes we know the results

So would there be an equivalent based on the number of "iterations" instead?

Upvotes: 1

Views: 493

Answers (1)

Jean-Francois T.
Jean-Francois T.

Reputation: 12940

I came up with the following function:

def first_elements(iterable, n:int):
    """Iterates over n elements of an iterable"""
    for _ in range(n):
        yield next(iterable)

And you could get a list as follows: first_10 = list(first_elements(long_generator, 10))

Is there some built-in or better/more elegant way?

Upvotes: 2

Related Questions