如何从迭代器/生成器中取出第一个元素并将其放回 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)

输出

first color is red
red
green
blue

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

在这里，您将打印循环x中的和的返回值（如果您调用它）。for.send()

0
1
2
3
4
5
5 # 5 出现两次，因为它被推迟了
6
7
8
9

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

# 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 # 推回 (num_pushed = 0)
5 # 推回 (num_pushed = 1)
5 # 推回 (num_pushed = 2)
5 # 不推回
6
7
8
9