Generadores y yield en Python

El post anterior terminó con veinte líneas de clase para contar hacia atrás desde tres. Aquí está lo mismo.

Tres líneas

def countdown(n):
    while n > 0:
        yield n
        n -= 1

print(list(countdown(3)))

for i in countdown(3):
    print(i)

Imprime:

[3, 2, 1]
3
2
1

Cualquier función que tenga un yield adentro es una función generadora. Llamarla no ejecuta el cuerpo: devuelve un generador, y un generador es un iterador. Python escribe __iter__ y __next__ por ti.

Nada corre hasta que lo pides

Esta es la parte que vale la pena mirar de cerca.

def noisy():
    print('  starting')
    yield 1
    print('  between')
    yield 2
    print('  finishing')

g = noisy()
print('created — the body has not run')
print('got', next(g))
print('got', next(g))
try:
    next(g)
except StopIteration:
    print('done')

Imprime:

created — the body has not run
  starting
got 1
  between
got 2
  finishing
done

Lee el orden. starting aparece después de created, no antes. El cuerpo corre solo cuando algo llama a next, y se detiene otra vez en cuanto llega a un yield.

return termina una función y descarta todo lo que era local a ella. yield la suspende con sus variables intactas, y el siguiente next() retoma en la línea siguiente. Cuando el cuerpo por fin termina, Python lanza StopIteration, la misma señal que el post anterior atrapaba a mano.

Por qué la pereza vale algo

import sys

squares_list = [n * n for n in range(1_000_000)]
squares_gen  = (n * n for n in range(1_000_000))

print('list:', sys.getsizeof(squares_list), 'bytes')
print('gen: ', sys.getsizeof(squares_gen), 'bytes')
print('same total:', sum(squares_list) == sum(squares_gen))

Imprime:

list: 8448728 bytes
gen:  200 bytes
same total: True

La lista guarda un millón de números. El generador guarda una función pausada. Los dos suman el mismo valor, y solo uno de ellos tuvo que caber en memoria.

Paréntesis redondos en lugar de corchetes: esa es toda la diferencia. Eso es una expresión generadora, una comprensión de listas que nunca construye la lista.

nums = [1, 2, 3, 4, 5, 6]

print(sum(n * n for n in nums if n % 2))
print(sorted((n * n for n in nums), reverse=True)[:3])

Imprime:

35
[36, 25, 16]

Cuando la expresión generadora es el único argumento, puedes quitarle los paréntesis: por eso sum(n * n for n in nums) se lee como se lee.

Ahora lo infinito está permitido

Una lista de todos los números naturales no es algo que puedas tener. Un generador de ellos sí, porque solo guarda el que está tratando en ese momento.

def naturals():
    n = 1
    while True:
        yield n
        n += 1

out = []
for n in naturals():
    if n * n > 500:
        break
    out.append(n * n)

print(out)

Imprime:

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484]

El while True nunca termina. El break es lo que lo termina, y el generador simplemente queda pausado, para siempre, y se recolecta.

Etapas que se jalan entre sí

Los generadores se encadenan. Cada etapa le pide a la anterior un solo elemento.

lines = ['17,ada,engineer', '', '18,grace,admiral', 'not a row', '19,alan,logician']

good   = (l for l in lines if l.count(',') == 2)
fields = (l.split(',') for l in good)
people = ((int(i), name) for i, name, _ in fields)

for pid, name in people:
    print(pid, name)

Imprime:

17 ada
18 grace
19 alan

No pasó nada hasta que el bucle for lo pidió. Después una línea recorrió las tres etapas, se imprimió, y siguió la próxima. En ningún momento existió una lista intermedia completa. Cambia lines por un archivo de 4 GB y el código no cambia.

yield from

Cuando el trabajo de un generador es pasar adelante lo que produce otro, yield from lo dice en una línea.

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

print(list(flatten([1, [2, [3, 4]], 5, [[6]]])))

Imprime:

[1, 2, 3, 4, 5, 6]

Sin eso escribirías un bucle que vuelve a hacer yield de cada elemento, lo cual funciona y se lee peor.

Las dos cosas que resignas

Un generador es un iterador, así que corre una sola vez: todo lo que dijo el post anterior aplica aquí.

g = (n for n in range(4))
print(list(g))
print(list(g))

Imprime:

[0, 1, 2, 3]
[]

Y no tiene longitud ni indexación, porque no sabe lo que todavía no produjo.

g = (n for n in range(5))
try:
    print(len(g))
except TypeError as err:
    print('TypeError:', err)

try:
    print(g[0])
except TypeError as err:
    print('TypeError:', err)

Imprime:

TypeError: object of type 'generator' has no len()
TypeError: 'generator' object is not subscriptable

Si necesitas la longitud o un índice, lo que necesitabas era una lista. Llama a list() sobre él y paga la memoria.

Qué recordar

• Un yield en el cuerpo de una función la convierte en función generadora. Llamarla no ejecuta nada.

• Cada next() corre el cuerpo hasta el siguiente yield y lo congela ahí, con variables locales y todo.

• Los paréntesis redondos te dan una expresión generadora: la comprensión que ya escribes, sin la lista.

• Los generadores pueden ser infinitos, y se encadenan en tuberías que nunca guardan más de un elemento por etapa.

• Una pasada, sin len, sin indexación. Ese es el precio, y casi siempre vale la pena pagarlo.