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

bodacydo picture bodacydo · Mar 11, 2010 · Viewed 36.2k times · Source

I can't figure out how to look ahead one element in a Python generator. As soon as I look it's gone.

Here is what I mean:

gen = iter([1,2,3])
next_value = gen.next()  # okay, I looked forward and see that next_value = 1
# but now:
list(gen)  # is [2, 3]  -- the first value is gone!

Here is a more real example:

gen = element_generator()
if gen.next_value() == 'STOP':
  quit_application()
else:
  process(gen.next())

Can anyone help me write a generator that you can look one element forward?

Answer

David Z picture David Z · Dec 30, 2014

For sake of completeness, the more-itertools package (which should probably be part of any Python programmer's toolbox) includes a peekable wrapper that implements this behavior. As the code example in the documentation shows:

>>> p = peekable(['a', 'b'])
>>> p.peek()
'a'
>>> next(p)
'a'

However, it's often possible to rewrite code that would use this functionality so that it doesn't actually need it. For example, your realistic code sample from the question could be written like this:

gen = element_generator()
command = gen.next_value()
if command == 'STOP':
  quit_application()
else:
  process(command)

(reader's note: I've preserved the syntax in the example from the question as of when I'm writing this, even though it refers to an outdated version of Python)