Generatory w Pythonie: yield, leniwa iteracja i wyrażenia generatorowe

Funkcja, która robi pauzy

Generator wygląda jak zwykła funkcja, ale zamiast obliczyć cały wynik i go zwrócić, oddaje (yield) jedną wartość naraz i zatrzymuje się między kolejnymi wartościami, dopóki ktoś nie poprosi o następną.

Najprostszy możliwy generator:

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

for number in count_up_to(5):
    print(number)

Zwróć uwagę na yield zamiast return. Gdy for po raz pierwszy prosi o wartość, Python wykonuje ciało funkcji, aż dojdzie do yield 1. Funkcja zatrzymuje się właśnie tam, oddaje pętli 1 i dokładnie pamięta, gdzie przerwała, łącznie ze zmiennymi. Następna iteracja wznawia działanie od tego miejsca: current += 1, powrót do while, yield 2. I tak dalej, aż warunek pętli przestanie być spełniony, a wtedy generator po prostu się kończy.

Na tym zatrzymywaniu i wznawianiu polega cała sztuczka.

Dlaczego nie zbudować po prostu listy?

Bo wersja z listą rezerwuje wszystkie wartości z góry:

def count_up_to_list(n):
    return list(range(1, n + 1))

print(count_up_to_list(5))

Przy 5 elementach to w porządku. Teraz wyobraź sobie, że chcesz 50 milionów liczb całkowitych, a interesuje cię tylko pierwsza spełniająca jakiś warunek. Wersja z listą tworzy 50 milionów liczb, z których większość i tak wyrzucisz. Wersja z generatorem tworzy dokładnie tyle, ile zużyje wywołujący. Gdy pętla for znajdzie to, czego szuka, i wykona break, generator po prostu przestaje działać.

To wzorzec, który warto sobie przyswoić: generatory pozwalają pisać kod iteracji bez decydowania z góry, jak dużej części wyniku będziesz potrzebować.

Wyrażenia generatorowe

Jeśli znasz już list comprehension, znasz też składnię: zamień nawiasy kwadratowe na okrągłe:

# List comprehension: rezerwuje cały wynik.
squares_list = [n * n for n in range(10)]

# Wyrażenie generatorowe: oblicza na żądanie.
squares_gen = (n * n for n in range(10))

print(squares_list)
print(squares_gen)

for s in squares_gen:
    print(s)

squares_gen jeszcze niczego nie oblicza. To tylko przepis. Iteracja wykonuje go krok po kroku.

Wyrażenia generatorowe świetnie sprawdzają się jako argumenty funkcji, które zużywają obiekt iterowalny:

numbers = [3, 8, 1, 5, 4, 9, 2]

print(sum(n * n for n in numbers))
print(max(n for n in numbers if n % 2 == 0))
print(any(n > 7 for n in numbers))

Bez pośredniej listy. sum, max i any czytają wartości po jednej, a właśnie tego potrzebują.

Odczyt dużego pliku linia po linii

To klasyczny praktyczny przypadek dla generatorów: przetwarzanie pliku, który jest za duży, by wczytać go do pamięci:

def parse_log_lines(path):
    with open(path) as f:
        for line in f:
            if line.startswith("ERROR"):
                yield line.rstrip()

for error in parse_log_lines("app.log"):
    print(error)

Plik jest czytany leniwie. Każde pobranie z generatora ściąga z dysku jedną linię, filtruje ją i oddaje. Zużycie pamięci pozostaje stałe niezależnie od rozmiaru pliku.

Raz i koniec

Po generatorze można przejść tylko raz. Gdy dojdziesz do końca, jest wyczerpany:

def first_three():
    yield 1
    yield 2
    yield 3

gen = first_three()

for n in gen:
    print("first time:", n)

for n in gen:
    print("second time:", n)

Druga pętla nic nie wypisuje. W generatorze nic już nie zostało.

Jeśli musisz iterować więcej niż raz, wywołaj funkcję generatora ponownie, żeby dostać nowy generator, albo zamień sekwencję w listę przez list(...) i iteruj po liście wielokrotnie. Wybierz według kosztu: ponowne budowanie jest w porządku, gdy praca jest tania, a lista, gdy sekwencja jest mała.

next() i ręczna iteracja

Nie musisz używać pętli for. next() pobiera jedną wartość naraz:

def ticks():
    yield "first"
    yield "second"
    yield "third"

gen = ticks()

print(next(gen))
print(next(gen))
print(next(gen))

# Czwarte wywołanie zgłosiłoby StopIteration.

StopIteration to sposób, w jaki generator sygnalizuje "to już koniec". Pętle for łapią go po cichu. W ręcznym kodzie możesz przekazać wartość domyślną do next(gen, default), żeby uniknąć wyjątku.

Generatory nieskończone

Skoro wartości powstają na żądanie, generator może reprezentować sekwencję bez końca, o ile konsument w pewnym momencie przestanie prosić:

def naturals():
    n = 1
    while True:
        yield n
        n += 1

for n in naturals():
    if n > 5:
        break
    print(n)

while True z yield w środku nie zawiesza programu. Znaczy tylko "dopóki ktoś prosi, wytwarzaj dalej". To konsument decyduje, kiedy przestać.

Ten wzorzec pojawia się przy strumieniach danych, pętlach zdarzeń i wszędzie tam, gdzie pobierasz wartości ze źródła bez określonej długości.

yield from: przekazanie do innego obiektu iterowalnego

Jeśli generator ma oddać każdą wartość z innego obiektu iterowalnego, yield from robi to w jednej linii:

def chain(*iterables):
    for it in iterables:
        yield from it

combined = chain([1, 2, 3], (4, 5, 6), "ab")
for x in combined:
    print(x)

Bez yield from trzeba by napisać zagnieżdżoną pętlę for z yield x w środku. Poprawnie przekazuje też wywołania send() i throw(), jeśli kiedyś ich użyjesz, ale w codziennym kodzie myśl o nim jak o "oddaj każdą wartość z tego obiektu".

Kiedy sięgnąć po generator

Trzy sygnały, że generator to właściwe narzędzie:

1. Sekwencja jest duża, być może nieskończona, albo jej pełne wytworzenie jest kosztowne.
2. Konsument może przestać przed końcem (na przykład break przy pierwszym trafieniu).
3. Chcesz łączyć przekształcenia (filtrowanie, mapowanie, branie pierwszych elementów) bez budowania pośrednich list.

A kiedy nie:

• Potrzebujesz dostępu swobodnego (seq[42]). Generatory idą tylko do przodu.
• Musisz iterować po tej samej sekwencji kilka razy. Użyj listy.
• Sekwencja jest mała i już ją masz. List comprehension będzie prostsze.

Generatory, list comprehensions i zwykłe listy to właściwe odpowiedzi w różnych zadaniach. Cała umiejętność polega na tym, by wybierać bez długiego zastanawiania się, a najszybciej wyrobisz ten instynkt, zauważając przy każdej pisanej iteracji, czy lepiej pasuje "wytwórz najpierw wszystko", czy "wytwarzaj po jednym".

Dalej: menedżery kontekstu w szczegółach

Znasz już większość idiomów, których Python używa do iteracji. Następne są menedżery kontekstu, czyli instrukcja with, które dobrze łączą się z generatorami przy strumieniowym pobieraniu danych z plików i połączeń sieciowych.