如何从iterator/generator获取第一个元素并将其放回Python中？

Question

我想从迭代器中取出第一个元素，分析它，然后把它放回原处，像没碰过一样和迭代器一起工作。在

现在我写道：

def prepend_iterator(element, it):
    yield element
    for element in it:
        yield element


def peek_first(it):
    first_element = next(it)
    it = prepend_iterator(first_element, it)
    return first_element, it

first_element, it = peek_first(it)

analyse(first_element)

continue_work(it)

有可能写得更好/更短？在

Answer 1

这里有一个例子itertools.tee

import itertools
def colors():
    for color in ['red', 'green', 'blue']:
        yield color

rgb = colors()
foo, bar = itertools.tee(rgb, 2)

#analize first element
first =  next(foo)
print('first color is {}'.format(first))

# consume second tee
for color in bar:
    print(color)

输出

^{pr2}$

Answer 2

注意，只有当您将非None值往后推时，这才有效。

如果您实现生成器函数（这就是您所拥有的）以便关心yield的返回值，那么您可以在生成器上“向后推”（使用.send()）：

# Generator
def gen():
    for val in range(10):
        while True:
            val = yield val
            if val is None: break

# Calling code
pushed = false
f = gen()
for x in f:
    print(x)
    if x == 5:
        print(f.send(5))
        pushed = True

在这里，您将打印来自for循环的x和{}（如果您调用它）的返回值。在

^{pr2}$

这只会让你后退一次。如果你想往后推更多次，你可以这样做：

# Generator
def gen():
    for val in range(10):
        while True:
            val = yield val
            if val is None: break

# Generator Wrapper
class Pushable:
    def __init__(self, g):
        self.g = g
        self._next = None

    def send(self, x):
        if self._next is not None:
            raise RuntimeError("Can't pushback twice without pulling")
        self._next = self.g.send(x)

    def __iter__(self):
        try:
            while True:
                # Have to clear self._next before yielding
                if self._next is not None:
                    (tmp, self._next) = (self._next, None)
                    yield tmp
                else:
                    yield next(self.g)

        except StopIteration: return

# Calling code
num_pushed = 0
f = Pushable(gen())
for x in f:
    print(x)
    if (x == 5) and (num_pushed in [0,1,2]):
        f.send(x)
        num_pushed += 1

产生：

0
1
2
3
4
5  # Pushed back (num_pushed = 0)
5  # Pushed back (num_pushed = 1)
5  # Pushed back (num_pushed = 2)
5  # Not pushed back
6
7
8
9