2.31. Iterátorok és generátorok

A for ciklus több munkát végzett, mint amennyinek látszik. Ez az oldal az iterátor protokollt tárgyalja, amelyen fut, valamint a yield kulcsszót, amely lehetővé teszi, hogy saját iterátorokat építs.

2.31.1. Az iterátor protokoll

Minden objektum, amelyen végig lehet iterálni, két metódust valósít meg:

  • __iter__() – visszaad egy iterátort az objektum elemei felett.

  • __next__() – az iterátoron visszaadja a következő elemet, vagy StopIteration kivételt vált ki, ha nincs több.

Az iter() beépített függvény az __iter__ metódust hívja; a next() az __next__ metódust hívja. Lépkedj végig egy listán kézzel:

it = iter([10, 20, 30])
print(next(it))    # 10
print(next(it))    # 20
print(next(it))    # 30
print(next(it))    # raises StopIteration
Egy for ciklus egyszer hívja meg az __iter__ metódust, hogy kapjon egy iterátort, majd ismételten hívja az __next__ metódust, amíg a StopIteration véget nem vet a ciklusnak.

A for voltaképpen a következő rövidített formája: „hívd meg egyszer az __iter__ metódust, majd iterálj az __next__ metóduson, amíg StopIteration nem jön.”

Mit csinál valójában a for x in items::

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

Minden lista, tuple, string, dict, set, fájlobjektum és generátor már megvalósítja az __iter__ és __next__ metódusokat – ezért működnek mind a for ciklussal.

2.31.2. A yield és a generátor függvények

Egy függvény, amely tartalmaz egy yield utasítást, egy generátor függvény. A meghívása nem futtatja le a törzset; egy generátor objektumot (egy iterátort) ad vissza, amely a törzset egyszerre egy yield erejéig futtatja:

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

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

Kimenet:

0
1
2

A next() minden hívása a következő yield pontig folytatja a függvényt, átadja azt az értéket a hívónak, és ott megáll. A lokális állapot (ebben az esetben az i) megmarad a folytatások között.

Szekvenciadiagram két életvonallal (a hívó és a generátor törzse). A hívó meghívja a count_up_to(3) függvényt, amely létrehoz egy generátort a törzs futtatása nélkül. Minden ezt követő next() a következő yield pontig futtatja a törzset, visszaadja az átadott értéket, és megáll. A negyedik next() leesik a végéről, és StopIteration kivételt vált ki. Az i változó megmarad a megállások között.

A next() a következő yield pontig futtatja a törzset, visszaadja az értéket, és megáll. A lokális állapot túléli a megállást.

A generátorok a legegyszerűbb módja egy szekvencia lusta előállításának – nem épül lista, az elemek csak akkor kerülnek kiszámításra, amikor a fogyasztó kéri őket, és a függvény akár örökké is adhat elemeket, ha akar.

2.31.3. Lusta folyamatok (pipeline)

A generátorok jól kombinálhatók. Az egyik generátor kimenete betáplálható egy másikba:

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)

Az értékek egyesével áramlanak végig a folyamaton – nincs köztes lista, nincs felső korlát beépítve a numbers függvénybe, és a fogyasztó (for v in pipeline) dönti el, mikor álljon le.

Három doboz balról jobbra: numbers(), squares(source), és a for-v-in-pipeline fogyasztó. Alattuk három ciklus van ábrázolva. Minden ciklusban a fogyasztó egy lekérési kérelmet küld balra a squares felé, amely egy lekérést küld balra a numbers felé; a numbers egy értéket ad át jobbra a squares felé, amely a négyzetre emelt értékét adja át jobbra a fogyasztónak.

A fogyasztón végrehajtott minden next() egy lekérést indít végig a láncon; az értékek csak akkor léteznek, amikor valami kéri őket.

2.31.3.1. yield from

Egy ciklus, amely elemeket húz egy másik iterálható objektumból, és mindegyiket átadja, elég gyakori ahhoz, hogy a Python rövidítést biztosítson rá. A yield from iter kifejezés sorrendben átadja az iterálható objektum által előállított minden értéket – mintha a generátornak egy for x in iter: yield x ciklusa lenne beágyazva:

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

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

Kimenet:

1
2
3
4
5
a
b
c

A yield from pontosan egyenértékű az explicit for ciklussal, csak rövidebb, és tisztán továbbítja a StopIteration kivételt a belső iterálható objektumtól a külső generátorig – hasznos, amikor több generátort fűzöl össze egymás után.

2.31.3.2. Amikor a yield kifogy

Egy generátor függvény végéről való leesés (vagy egy explicit return elérése) automatikusan StopIteration kivételt vált ki. Nincs szükség kézzel kiváltani; a körülvevő for ciklus észleli, és véget ér.

Használj generátorokat, ha az előállító kód természetes módon ciklusként van megírva néhány yield ponttal; használj egyszerű listakifejezést (list comprehension), ha valóban szükséged van az egész szekvenciára a memóriában.