Как взять первый элемент из итератора / генератора и вернуть его обратно в Python? - PullRequest
Новая
       8

Как взять первый элемент из итератора / генератора и вернуть его обратно в Python?

0 голосов
/ 10 декабря 2018

Я хотел бы взять первый элемент из итератора, проанализировать его, затем поместить обратно и работать с итератором, как если бы он не был затронут.

Пока я написал: 

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)

можно написать лучше / короче?

Ответы [ 2 ]

0 голосов
/ 10 декабря 2018

Вот пример с 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
0 голосов
/ 10 декабря 2018

Обратите внимание, что это будет работать только в том случае, если вы отбрасываете не-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 appears twice because it was pushed back
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  # 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
...