How to Look Ahead One Element (Peek) in a Python Generator

How to look ahead one element (peek) in a Python generator?

The Python generator API is one way: You can't push back elements you've read. But you can create a new iterator using the itertools module and prepend the element:

import itertools

gen = iter([1,2,3])
peek = gen.next()
print list(itertools.chain([peek], gen))

Using lookahead with generators

You can write a wrapper that buffers some number of items from the generator, and provides a lookahead() function to peek at those buffered items:

class Lookahead:
    def __init__(self, iter):
        self.iter = iter
        self.buffer = []

    def __iter__(self):
        return self

    def next(self):
        if self.buffer:
            return self.buffer.pop(0)
        else:
            return self.iter.next()

    def lookahead(self, n):
        """Return an item n entries ahead in the iteration."""
        while n >= len(self.buffer):
            try:
                self.buffer.append(self.iter.next())
            except StopIteration:
                return None
        return self.buffer[n]

How to look ahead one element (peek) in a Python generator?

The Python generator API is one way: You can't push back elements you've read. But you can create a new iterator using the itertools module and prepend the element:

import itertools

gen = iter([1,2,3])
peek = gen.next()
print list(itertools.chain([peek], gen))

Zipped Python generators with 2nd one being shorter: how to retrieve element that is silently consumed

If you want to reuse code, the easiest solution is:

from more_itertools import peekable

a = peekable(a)
b = peekable(b)

while True:
    try:
        a.peek()
        b.peek()
    except StopIteration:
        break
    x = next(a)
    y = next(b)
    print(x, y)

print(list(a), list(b))  # Misses nothing.

You can test this code out using your setup:

def my_gen(n: int):
    yield from range(n)

a = my_gen(10)
b = my_gen(8)

It will print:

0 0
1 1
2 2
3 3
4 4
5 5
6 6
7 7
[8, 9] []

Elegant way to fetch an element from a generator without actually using it?

It is not possible to get an element from a generator without advancing the generator. That's how generators work. You can use tricks to store the element to use it again later (which you seem to already know about). You could also write a class providing that behavior that wraps a generator with a peek method. But you have to do all that external to the generator mechanism. There isn't any way to make the generator itself not go forward when you get an element.

---

Related Topics