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