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
A for loop calls __iter__ once to get an iterator, then calls __next__ repeatedly until StopIteration ends the loop.

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.

Sequence diagram with two lifelines (caller and generator body). The caller calls count_up_to(3), which creates a generator without running the body. Each subsequent next() runs the body until the next yield, returns the yielded value, and pauses. The fourth next() falls off the end and raises StopIteration. The variable i is preserved across pauses.

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ć.

Three boxes left-to-right: numbers(), squares(source), and the for-v-in-pipeline consumer. Three cycles are drawn beneath. In each cycle, the consumer sends a pull request leftward to squares, which sends a pull leftward to numbers; numbers yields a value rightward to squares, which yields its squared value rightward to the consumer.

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.