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

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.

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

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.

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.