I just got bit by abusing StopIteration
within some nested generators (using CPython 3.6.9), hadn't enabled PEP 479 (from __future__ import generator_stop
), and had some bad hacky code that used next(iter(iterable))
that prematurely signaled a stop.
While PEP 479 would catch StopIteration
from bubbling out of generators, I think I would still run into this in nested for-loops.
For now, I'm going to replace any usage of next(iter(...))
with the following:
def take(iterable, *, n):
"""
Robustly gets the first n items from an iterable and returns them as a
list.
You should always use this function in lieu of `next(iter(...))`! e.g.
instead of:
my_first = next(iter(container))
you should instead do:
my_first, = take(container, n=1)
Throws RuntimeError if the iterable cannot yield n items.
"""
iterator = iter(iterable)
out = []
for _ in range(n):
try:
out.append(next(iterator))
except StopIteration:
raise RuntimeError("Premature StopIteration encountered!")
return out
My question is: Is function like this already in the stdlib for Python?
I checked through python.org
s latest docs (for 3.9) in itertools
and builtins
, and the closest thing I could see was takewhile
, but meh on that. I could also convert to a list
or any other indexable container, but I'd like to avoid needing to iterate over everything just to access the first thing.
itertools.islice
does this (and more), without converting to a list or erroring out if not enough elements are produced.
You could write your function in terms of this one cleanly:
def take(iterable, *, n):
li = list(itertools.islice(iterable, n))
if len(li) != n:
raise RuntimeError("too short iterable for take")
return li