2.31. Iteratorer och generatorer

for-loopen har gjort mer arbete än det ser ut som. Den här sidan täcker iteratorprotokollet som den bygger på, och nyckelordet yield som låter dig bygga dina egna iteratorer.

2.31.1. Iteratorprotokollet

Varje objekt som går att loopa över implementerar två metoder:

  • __iter__() – returnerar en iterator över objektets element.

  • __next__() – på iteratorn, returnerar nästa element eller höjer StopIteration när det inte finns fler.

Den inbyggda iter() anropar __iter__; next() anropar __next__. Stega igenom en lista för hand:

it = iter([10, 20, 30])
print(next(it))    # 10
print(next(it))    # 20
print(next(it))    # 30
print(next(it))    # raises StopIteration
En for-loop anropar __iter__ en gång för att få en iterator, och anropar sedan __next__ upprepade gånger tills StopIteration avslutar loopen.

for är socker för ”anropa __iter__ en gång, loopa sedan på __next__ tills StopIteration.”

Vad for x in items: faktiskt gör:

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

Varje lista, tupel, sträng, dict, mängd, filobjekt och generator implementerar redan __iter__ och __next__ – vilket är anledningen till att de alla fungerar med for.

2.31.2. yield och generatorfunktioner

En funktion som innehåller en yield-sats är en generatorfunktion. Att anropa den kör inte kroppen; den returnerar ett generatorobjekt (en iterator) som kör kroppen ett yield i taget:

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

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

Utdata:

0
1
2

Varje anrop till next() återupptar funktionen fram till nästa yield, lämnar det värdet till anroparen och pausar där. Det lokala tillståndet (i i detta fall) bevaras mellan återupptagningarna.

Sekvensdiagram med två livslinjer (anropare och generatorkropp). Anroparen anropar count_up_to(3), vilket skapar en generator utan att köra kroppen. Varje efterföljande next() kör kroppen fram till nästa yield, returnerar det lämnade värdet och pausar. Det fjärde next() faller av slutet och höjer StopIteration. Variabeln i bevaras över pauserna.

next() kör kroppen fram till nästa yield, lämnar tillbaka värdet och pausar. Det lokala tillståndet överlever pausen.

Generatorer är det enklaste sättet att producera en sekvens lat – ingen lista byggs, element beräknas endast när konsumenten begär dem, och funktionen kan lämna element i all evighet om den vill.

2.31.3. Lata pipelines

Generatorer komponeras väl. En generators utdata kan mata en annan:

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)

Värdena flödar genom pipelinen ett i taget – ingen mellanliggande lista, ingen övre gräns inbyggd i numbers, och konsumenten (for v in pipeline) bestämmer när det ska sluta.

Tre rutor från vänster till höger: numbers(), squares(source) och konsumenten for-v-in-pipeline. Tre cykler är ritade nedanför. I varje cykel skickar konsumenten en begäran (pull) åt vänster till squares, som skickar en begäran åt vänster till numbers; numbers lämnar ett värde åt höger till squares, som lämnar sitt kvadrerade värde åt höger till konsumenten.

Varje next() på konsumenten utlöser en begäran genom kedjan; värden existerar endast när något ber om dem.

2.31.3.1. yield from

En loop som drar element från en annan iterabel och lämnar varje ett är tillräckligt vanligt för att Python tillhandahåller en genväg. Uttrycket yield from iter lämnar varje värde som den iterabla producerar, i ordning – som om generatorn hade en for x in iter: yield x-loop inlagd:

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

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

Utdata:

1
2
3
4
5
a
b
c

yield from är exakt likvärdigt med den explicita for-loopen, bara kortare, och det propagerar StopIteration från den inre iterabla upp till den yttre generatorn på ett rent sätt – användbart när flera generatorer kedjas ihop från början till slut.

2.31.3.2. När yield tar slut

Att falla av slutet på en generatorfunktion (eller att träffa en explicit return) höjer StopIteration automatiskt. Det finns inget behov av att höja den för hand; den omgivande for-loopen ser den och avslutas.

Använd generatorer när den producerande koden naturligt skrivs som en loop med några yield-punkter; använd en vanlig listomfattning när du verkligen behöver hela sekvensen i minnet.