CODEWITHSUNDEEP
pythonpython-3.xiterator
porton
porton

Reputation: 5805

Excluding the first element from an iterator (Python)

I have a generator function A.

For example (in reality I have a more complex function A),

def A():
    yield from [i**2 for i in range(20)]

Writing another generator function B, I want to enumerate all elements that A returns except the first element.

What are concise ways to implement this in Python 3?

Upvotes: 2

Views: 2290

Answers (3)

abarnert
abarnert

Reputation: 365657

Usually, you don't need this in an expression, so you just call next(it), ignoring the results, to consume and discard the first element.

However, if the iterator might be empty, you have to decide what you want to happen:

  • Maybe you want to raise StopIteration, in which case next(it) is fine.
  • Maybe you want to raise something else, in which case you next(it) inside an except StopIteration: raise SomethingElse().
  • Maybe you just want to leave the iterator empty, in which case you can call next(it, None).

You can find examples of these in the stdlib and docs. For example, if you scan through the recipes in itertools:

def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

This is doing exactly what you want to do—skip the first element of b. And if iterable is empty, you don't want an error here; you just want to iterate nothing. So, next(b, None).

What if you need to do this in the middle of an expression?

Then you can write a function that skips the first element:

def skip_first(iterable):
    it = iter(iterable)
    next(it, None)
    return it

(Again, you have to decide what you want to happen for an empty iterable.)

This returns a first-skipped version of the iterator, so you can use it inline. (It also mutates the iterator you passed in, of course, but you normally only use on a temporary value that you're not keeping any references to, so that's not a problem.)

Or, if you need to return a generator instead of an arbitrary iterator (usually you don't):

def skip_first(iterable):
    it = iter(iterable)
    next(it, None)
    yield from it

Or you can use the more general version of the same idea, itertools.islice. The following have the same effect:

it = skip_first(it)
it = itertools.islice(it, 1, None)

While we're on the itertools recipes, it's worth looking at consume:

def consume(iterator, n=None):
    "Advance the iterator n-steps ahead. If n is None, consume entirely."
    # Use functions that consume iterators at C speed.
    if n is None:
        # feed the entire iterator into a zero-length deque
        collections.deque(iterator, maxlen=0)
    else:
        # advance to the empty slice starting at position n
        next(islice(iterator, n, n), None)

Forget the None part; the interesting bit is that it skips n elements with an islice and a next. (Notice that it's mutating iterator in-place, not returning something.)

Upvotes: 4

rafaelc
rafaelc

Reputation: 59274

Disclaimer: See @abarnert and @Solaxun's answers above.

Just thought that the following should be mentioned

If you have e.g. original = iter((1,2,3,4,5))

Then

first, remaining = next(original), original

where remaining is the iterator without the first element.

Upvotes: 2

Solaxun
Solaxun

Reputation: 2792

Use itertools.islice:

itertools.islice(generator,1,None)

Upvotes: 11

Related Questions