6

What is the most efficient way to get a fixed number of items from a generator?

What I currently do is using zip and range. In this example I take chunks of size 3 from the generator.

def f():
  x = 0
  while x < 21:
    yield x
    x += 1

g = f()

while True:
  x = [i for _, i in zip(range(3), g)]
  if not x:
    break
  print x

The background is that the database I use provides a generator object for query results. Than I fill a fixed size numpy array with data and process it as one batch.

tauran
  • 7,986
  • 6
  • 41
  • 48

1 Answers1

6

Use itertools.islice:

import itertools

for elem in itertools.islice(f(), 3):
    print elem

and directly into your numpy array:

my_arr = np.array(itertools.islice(f(), 3))
eumiro
  • 207,213
  • 34
  • 299
  • 261