2.31. Iteradores y generadores

El bucle for ha estado haciendo más trabajo del que parece. Esta página cubre el protocolo de iteración en el que se basa, y la palabra clave yield que te permite construir tus propios iteradores.

2.31.1. El protocolo de iteración

Todo objeto sobre el que se puede iterar implementa dos métodos:

  • __iter__() – devuelve un iterador sobre los elementos del objeto.

  • __next__() – en el iterador, devuelve el siguiente elemento o lanza StopIteration cuando no quedan más.

La función incorporada iter() llama a __iter__; next() llama a __next__. Recorre una lista manualmente:

it = iter([10, 20, 30])
print(next(it))    # 10
print(next(it))    # 20
print(next(it))    # 30
print(next(it))    # raises StopIteration
Un bucle for llama a __iter__ una vez para obtener un iterador, luego llama a __next__ repetidamente hasta que StopIteration finaliza el bucle.

for es azúcar sintáctico para «llamar a __iter__ una vez, luego iterar sobre __next__ hasta StopIteration».

Lo que for x in items: hace en realidad:

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

Toda lista, tupla, cadena, diccionario, conjunto, objeto de archivo y generador ya implementa __iter__ y __next__ – por eso todos funcionan con for.

2.31.2. yield y las funciones generadoras

Una función que contiene una sentencia yield es una función generadora. Llamarla no ejecuta el cuerpo; devuelve un objeto generador (un iterador) que ejecuta el cuerpo un yield a la vez:

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

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

Salida:

0
1
2

Cada llamada a next() reanuda la función hasta el siguiente yield, entrega ese valor a quien la llamó y se pausa ahí. El estado local (i en este caso) se conserva entre reanudaciones.

Diagrama de secuencia con dos líneas de vida (el llamador y el cuerpo del generador). El llamador llama a count_up_to(3), lo que crea un generador sin ejecutar el cuerpo. Cada next() posterior ejecuta el cuerpo hasta el siguiente yield, devuelve el valor producido y se pausa. La cuarta llamada a next() se sale del final y lanza StopIteration. La variable i se conserva entre pausas.

next() ejecuta el cuerpo hasta el siguiente yield, devuelve el valor y se pausa. El estado local sobrevive a la pausa.

Los generadores son la forma más sencilla de producir una secuencia de manera perezosa – no se construye ninguna lista, los elementos se calculan solo cuando el consumidor los pide, y la función puede producir elementos indefinidamente si quiere.

2.31.3. Tuberías perezosas

Los generadores se componen bien. La salida de un generador puede alimentar a otro:

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)

Los valores fluyen por la tubería de uno en uno – sin lista intermedia, sin límite superior incorporado en numbers, y el consumidor (for v in pipeline) decide cuándo detenerse.

Tres cajas de izquierda a derecha: numbers(), squares(source), y el consumidor for-v-in-pipeline. Debajo se dibujan tres ciclos. En cada ciclo, el consumidor envía una solicitud de extracción hacia la izquierda a squares, que envía una extracción hacia la izquierda a numbers; numbers produce un valor hacia la derecha a squares, que produce su valor al cuadrado hacia la derecha al consumidor.

Cada next() en el consumidor desencadena una extracción a través de la cadena; los valores existen solo cuando algo los pide.

2.31.3.1. yield from

Un bucle que extrae elementos de otro iterable y produce cada uno es lo bastante común como para que Python ofrezca un atajo. La expresión yield from iter produce cada valor que genera el iterable, en orden – como si el generador tuviera un bucle for x in iter: yield x en línea:

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

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

Salida:

1
2
3
4
5
a
b
c

yield from es exactamente equivalente al bucle for explícito, solo que más corto, y propaga StopIteration desde el iterable interno hacia el generador externo de forma limpia – útil al encadenar varios generadores de extremo a extremo.

2.31.3.2. Cuando yield se agota

Salirse del final de una función generadora (o alcanzar un return explícito) lanza StopIteration automáticamente. No hay necesidad de lanzarla manualmente; el bucle for circundante la detecta y termina.

Usa generadores cuando el código productor se escribe de forma natural como un bucle con unos pocos puntos yield; usa una comprensión de lista normal cuando realmente necesites toda la secuencia en memoria.