2.31. Iteratory i generatory

Pętla for wykonuje więcej pracy, niż się wydaje. Ta strona omawia protokół iteratora, na którym ona działa, oraz słowo kluczowe yield, które pozwala budować własne iteratory.

2.31.1. Protokół iteratora

Każdy obiekt, po którym można iterować, implementuje dwie metody:

• __iter__() – zwraca iterator po elementach obiektu.

• __next__() – na iteratorze zwraca następny element lub zgłasza StopIteration, gdy nie ma już więcej elementów.

Wbudowana funkcja iter() wywołuje __iter__; next() wywołuje __next__. Przejdźmy przez listę ręcznie:

it = iter([10, 20, 30])
print(next(it))    # 10
print(next(it))    # 20
print(next(it))    # 30
print(next(it))    # raises StopIteration

for jest cukrem składniowym dla „wywołaj __iter__ raz, a następnie iteruj po __next__ aż do StopIteration„.

Co faktycznie robi for x in items::

_it = iter(items)
while True:
    try:
        x = next(_it)
    except StopIteration:
        break
    # ... loop body ...

Każda lista, krotka, łańcuch znaków, słownik, zbiór, obiekt pliku i generator już implementuje __iter__ oraz __next__ – dlatego wszystkie działają z for.

2.31.2. yield i funkcje generatorów

Funkcja, która zawiera instrukcję yield, jest funkcją generatora. Jej wywołanie nie uruchamia ciała; zwraca ono obiekt generatora (iterator), który uruchamia ciało po jednym yield na raz:

def count_up_to(n):
    i = 0
    while i < n:
        yield i
        i += 1

for value in count_up_to(3):
    print(value)

Wynik:

0
1
2

Każde wywołanie next() wznawia funkcję aż do następnego yield, przekazuje tę wartość wywołującemu i wstrzymuje się w tym miejscu. Stan lokalny (i w tym przypadku) jest zachowywany między wznowieniami.

next() uruchamia ciało aż do następnego yield, przekazuje wartość z powrotem i wstrzymuje się. Stan lokalny przetrwa wstrzymanie.

Generatory są najprostszym sposobem na leniwe wytwarzanie sekwencji – żadna lista nie jest budowana, elementy są obliczane tylko wtedy, gdy konsument o nie poprosi, a funkcja może generować elementy w nieskończoność, jeśli zechce.

2.31.3. Leniwe potoki

Generatory dobrze się komponują. Wyjście jednego generatora może zasilać kolejny:

def numbers():
    i = 0
    while True:
        yield i
        i += 1

def squares(source):
    for x in source:
        yield x * x

pipeline = squares(numbers())

for v in pipeline:
    if v > 100:
        break
    print(v)

Wartości przepływają przez potok jedna po drugiej – bez pośredniej listy, bez wbudowanego górnego ograniczenia w numbers oraz to konsument (for v in pipeline) decyduje, kiedy zatrzymać.

Każde next() na konsumencie wyzwala jedno pobranie przez łańcuch; wartości istnieją tylko wtedy, gdy coś o nie poprosi.

2.31.3.1. yield from

Pętla, która pobiera elementy z innego obiektu iterowalnego i generuje każdy z nich, jest na tyle powszechna, że Python udostępnia skrót. Wyrażenie yield from iter generuje każdą wartość wytwarzaną przez obiekt iterowalny, w kolejności – tak jakby generator miał wbudowaną pętlę for x in iter: yield x:

def chain(*sources):
    for source in sources:
        yield from source

for v in chain([1, 2, 3], (4, 5), "abc"):
    print(v)

Wynik:

1
2
3
4
5
a
b
c

yield from jest dokładnym odpowiednikiem jawnej pętli for, tylko krótszym, i czysto propaguje StopIteration z wewnętrznego obiektu iterowalnego do zewnętrznego generatora – przydatne przy łączeniu kilku generatorów jeden za drugim.

2.31.3.2. Gdy yield się wyczerpie

Wyjście poza koniec funkcji generatora (lub natrafienie na jawne return) automatycznie zgłasza StopIteration. Nie ma potrzeby zgłaszania go ręcznie; otaczająca pętla for widzi to i kończy działanie.

Używaj generatorów, gdy kod wytwarzający jest naturalnie zapisany jako pętla z kilkoma punktami yield; używaj zwykłego wyrażenia listowego, gdy naprawdę potrzebujesz całej sekwencji w pamięci.