如何在Python生成器中向前查看一个元素（peek）？

Question

我搞不清楚怎么在Python的生成器里向前看一个元素。一看就没了。

我说的意思是：

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!

这里有一个更真实的例子：

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

有没有人能帮我写一个可以向前看一个元素的生成器？

---

_{另见： 重置Python中的生成器对象}

Answer 1

好的，虽然我来得有点晚，但我看到这个问题，发现没有一个答案让我满意。于是我想出了这个元生成器：

class Peekorator(object):

    def __init__(self, generator):
        self.empty = False
        self.peek = None
        self.generator = generator
        try:
            self.peek = self.generator.next()
        except StopIteration:
            self.empty = True

    def __iter__(self):
        return self

    def next(self):
        """
        Return the self.peek element, or raise StopIteration
        if empty
        """
        if self.empty:
            raise StopIteration()
        to_return = self.peek
        try:
            self.peek = self.generator.next()
        except StopIteration:
            self.peek = None
            self.empty = True
        return to_return

def simple_iterator():
    for x in range(10):
        yield x*3

pkr = Peekorator(simple_iterator())
for i in pkr:
    print i, pkr.peek, pkr.empty

结果是：

0 3 False
3 6 False
6 9 False
9 12 False    
...
24 27 False
27 None False

也就是说，在你遍历列表的过程中，随时都可以访问到下一个项目。

Answer 2

为了完整性，more-itertools这个包（应该是每个Python程序员工具箱里必备的）里有一个叫peekable的功能，可以实现这种效果。就像文档中的代码示例所展示的那样：

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

---

不过，很多时候我们可以重写那些需要这个功能的代码，让它其实不需要这个功能。比如，你在问题中提到的代码示例可以这样写：

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

（读者注意：我保留了问题中示例的语法，尽管它提到的是一个过时的Python版本）

Answer 3

Python的生成器是一种特殊的工具，它的工作方式是你不能把已经读取过的元素再放回去。不过，你可以使用itertools模块来创建一个新的迭代器，并把元素放到最前面。

import itertools

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