How to take the first N items from a generator or list? [duplicate]

With linq I would

var top5 = array.Take(5);

How to do this with Python?

Solution 1:

Slicing a list

top5 = array[:5]

To slice a list, there's a simple syntax: array[start:stop:step]
You can omit any parameter. These are all valid: array[start:], array[:stop], array[::step]

Slicing a generator

import itertools
top5 = itertools.islice(my_list, 5) # grab the first five elements

You can't slice a generator directly in Python. itertools.islice() will wrap an object in a new slicing generator using the syntax itertools.islice(generator, start, stop, step)
Remember, slicing a generator will exhaust it partially. If you want to keep the entire generator intact, perhaps turn it into a tuple or list first, like: result = tuple(generator)

Solution 2:

import itertools

top5 = itertools.islice(array, 5)

Solution 3:

In my taste, it's also very concise to combine zip() with xrange(n) (or range(n) in Python3), which works nice on generators as well and seems to be more flexible for changes in general.

# Option #1: taking the first n elements as a list
[x for _, x in zip(xrange(n), generator)]

# Option #2, using 'next()' and taking care for 'StopIteration'
[next(generator) for _ in xrange(n)]

# Option #3: taking the first n elements as a new generator
(x for _, x in zip(xrange(n), generator))

# Option #4: yielding them by simply preparing a function
# (but take care for 'StopIteration')
def top_n(n, generator):
    for _ in xrange(n): yield next(generator)

Solution 4:

@Shaikovsky's answer is excellent (…and heavily edited since I posted this answer), but I wanted to clarify a couple of points.

[next(generator) for _ in range(n)]

This is the most simple approach, but throws StopIteration if the generator is prematurely exhausted.

On the other hand, the following approaches return up to n items which is preferable in many circumstances:

List: [x for _, x in zip(range(n), records)]

Generator: (x for _, x in zip(range(n), records))