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

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.

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() 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.

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.

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.