如何从迭代器/生成器中取出第一个元素并将其放回 Python？

Question

I would like to take first element from iterator, analyse it, then put it back and work with iterator as if it was not touched.

For now I wrote:

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)

it is possible to write better/shorter?

Answer 1

Here is an example wit 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)

output

first color is red
red
green
blue

Answer 2

Here I present a simple method that exploits concatenation of generators.

import itertools
def concat_generators(*args: Iterable[Generator]) -> Generator:
    r"""
    Concat generators by yielding from first, second, ..., n-th
    """
    for gen in args:
        yield from gen

your_generator = (i for i in range(10))
first_element = next(your_generator)

# then you could do this
your_generator = concat_generators([first_element], your_generator)
# or this
your_generator = itertools.chain([first_element], your_generator)

Answer 3

Note this will only work if you're pushing back non- None values.

If you implement your generator function (which is what you have) so that you care about the return value of yield , you can "push back" on the generator (with .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

Here, you're printing both the x from the for loop and the return value of .send() (if you call it).

0
1
2
3
4
5
5   # 5 appears twice because it was pushed back
6
7
8
9

This will only work let you push back once. If you want to push back more times than that, you do something like:

# 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

Produces:

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