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

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.

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

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.

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.