2.31. Iterátory a generátory

Smyčka for odvádí víc práce, než se zdá. Tato stránka popisuje protokol iterátoru, na kterém běží, a klíčové slovo yield, které vám umožňuje vytvářet vlastní iterátory.

2.31.1. Protokol iterátoru

Každý objekt, přes který lze iterovat, implementuje dvě metody:

• __iter__() – vrátí iterátor přes položky objektu.

• __next__() – na iterátoru vrátí další položku nebo vyvolá StopIteration, když už žádné nejsou.

Vestavěná funkce iter() volá __iter__; next() volá __next__. Projděte seznam ručně:

it = iter([10, 20, 30])
print(next(it))    # 10
print(next(it))    # 20
print(next(it))    # 30
print(next(it))    # raises StopIteration

for je syntaktická zkratka pro „zavolej jednou __iter__, pak iteruj přes __next__, dokud nenastane StopIteration.“

Co ve skutečnosti dělá for x in items::

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

Každý seznam, n-tice, řetězec, slovník, množina, objekt souboru a generátor už implementuje __iter__ a __next__ – proto všechny fungují s for.

2.31.2. yield a funkce generátoru

Funkce, která obsahuje příkaz yield, je funkce generátoru. Její zavolání nespustí tělo; vrátí objekt generátoru (iterátor), který tělo vykonává vždy po jednom yield:

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

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

Výstup:

0
1
2

Každé volání next() obnoví funkci až do dalšího yield, předá tuto hodnotu volajícímu a tam ji pozastaví. Lokální stav (v tomto případě i) je mezi obnoveními zachován.

next() spustí tělo až do dalšího yield, předá hodnotu zpět a pozastaví se. Lokální stav pozastavení přežije.

Generátory jsou nejjednodušší způsob, jak líně vytvářet sekvenci – nebuduje se žádný seznam, položky se počítají jen tehdy, když si je konzument vyžádá, a funkce může poskytovat položky donekonečna, pokud chce.

2.31.3. Líné pipeline

Generátory se dobře skládají. Výstup jednoho generátoru může napájet jiný:

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)

Hodnoty protékají pipeline po jedné – žádný mezilehlý seznam, žádná horní mez zabudovaná do numbers a konzument (for v in pipeline) rozhoduje, kdy se zastavit.

Každé next() na konzumentovi spustí jedno vytažení skrz řetězec; hodnoty existují jen tehdy, když si je něco vyžádá.

2.31.3.1. yield from

Smyčka, která táhne položky z jiné iterovatelné struktury a každou poskytuje, je natolik běžná, že Python poskytuje zkratku. Výraz yield from iter poskytne každou hodnotu, kterou iterovatelná struktura vyprodukuje, v pořadí – jako kdyby měl generátor vloženou smyčku 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)

Výstup:

1
2
3
4
5
a
b
c

yield from je přesně ekvivalentní explicitní smyčce for, jen kratší, a čistě propaguje StopIteration z vnitřní iterovatelné struktury nahoru do vnějšího generátoru – užitečné při řetězení několika generátorů za sebou.

2.31.3.2. Když yield dojde

Spadnutí za konec funkce generátoru (nebo dosažení explicitního return) automaticky vyvolá StopIteration. Není potřeba ji vyvolávat ručně; okolní smyčka for ji uvidí a skončí.

Použijte generátory, když je produkující kód přirozeně napsán jako smyčka s několika body yield; použijte prostou list comprehension, když skutečně potřebujete celou sekvenci v paměti.