How to take first element from iterator/generator and put it back in 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

EDIT (05 May 2022): To expand this answer, if you don't mind installing extra [third-party] package, there is more-itertools, that provide convenient tools to lookahead and lookback when working with iterator. These tools peek at an iterable’s values without advancing it.

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