CODEWITHSUNDEEP
pythonpython-itertools
Chris
Chris

Reputation: 1702

Dealing with `StopIteration` return from a next() call in Python

I'm using the below to skip a group of records when a certain condition is met:

if (condition met):
    ...
    [next(it) for x in range(19)]

Where it is an itertuples object created to speed up looping through a large dataframe (yes, the loop is necessary).

it = df.itertuples()
for row in it:
    ...

What's the idiomatic way of dealing with a StopIteration return from the next call (presumably due to reaching the end of the dataframe)?

Upvotes: 1

Views: 44

Answers (2)

blhsing
blhsing

Reputation: 107015

You can use itertools.islice with a stop of 19 to iterate at most 19 records from the iterator while safely stopping early if the iterator produces less than 19 records. To consume the iterator, you can construct a temporary list from the output of islice:

from itertools import islice

list(islice(it, 19))

If the space required by the temporary list is a concern, you can construct a collections.deque object with a maxlen of 0 instead:

from itertools import islice
from collections import deque

deque(islice(it, 19), maxlen=0)

Upvotes: 1

juanpa.arrivillaga
juanpa.arrivillaga

Reputation: 96257

So, in general, the solution to a potential exception being raised is to use exception handling. But in the case of next, you can simply use the second argument (which will make it return a default value in case the iterator is exhausted). But you shouldn't use a list comprehension like this anyway, so the most basic solution:

if condition:
    for _ in range(19):
        next(it, None)

But perhaps more elegantly (using the approach from the consume recipe in the itertools docs):

import itertools

def advance(it, n):
    next(itertools.islice(it, n, n), None)

...

if condition:
    advance(it, 19)

But again, a perfectly idiomatic solution would have been:

try:
    [next(it) for x in range(19)]
except StopIteration:
    pass

leaving aside the use of a list comprehension for side-effects

Upvotes: 2

Related Questions