pythongeneratorpython-3.6stopiteration

Is there a built-in `take(iterable, n)` function in Python3?


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.orgs 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.


Solution

  • 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