2.31. Iteradores e geradores

O laço for vem fazendo mais trabalho do que aparenta. Esta página cobre o protocolo de iterador sobre o qual ele opera e a palavra-chave yield, que permite construir seus próprios iteradores.

2.31.1. O protocolo de iterador

Todo objeto que pode ser percorrido em um laço implementa dois métodos:

• __iter__() – retorna um iterador sobre os itens do objeto.

• __next__() – no iterador, retorna o próximo item ou levanta StopIteration quando não há mais nenhum.

A função embutida iter() chama __iter__; next() chama __next__. Percorra uma lista manualmente:

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

for é açúcar sintático para “chamar __iter__ uma vez, depois iterar sobre __next__ até StopIteration.”

O que for x in items: realmente faz:

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

Toda lista, tupla, string, dict, set, objeto de arquivo e gerador já implementa __iter__ e __next__ – e é por isso que todos funcionam com for.

2.31.2. yield e funções geradoras

Uma função que contém uma instrução yield é uma função geradora. Chamá-la não executa o corpo; ela retorna um objeto gerador (um iterador) que executa o corpo um yield por vez:

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

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

Saída:

0
1
2

Cada chamada a next() retoma a função até o próximo yield, entrega esse valor ao chamador e pausa ali. O estado local (i neste caso) é preservado entre as retomadas.

next() executa o corpo até o próximo yield, devolve o valor e pausa. O estado local sobrevive à pausa.

Geradores são a maneira mais fácil de produzir uma sequência preguiçosamente – nenhuma lista é construída, os itens são calculados apenas quando o consumidor os solicita, e a função pode produzir itens indefinidamente, se quiser.

2.31.3. Pipelines preguiçosos

Geradores se compõem bem. A saída de um gerador pode alimentar outro:

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)

Os valores fluem pelo pipeline um de cada vez – sem lista intermediária, sem limite superior embutido em numbers, e o consumidor (for v in pipeline) decide quando parar.

Cada next() no consumidor dispara um puxar através da cadeia; os valores só existem quando algo os solicita.

2.31.3.1. yield from

Um laço que puxa itens de outro iterável e produz cada um é comum o suficiente para que o Python ofereça um atalho. A expressão yield from iter produz cada valor que o iterável gera, em ordem – como se o gerador tivesse um laço for x in iter: yield x embutido:

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

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

Saída:

1
2
3
4
5
a
b
c

yield from é exatamente equivalente ao laço for explícito, apenas mais curto, e propaga StopIteration do iterável interno para o gerador externo de forma limpa – útil ao encadear vários geradores de ponta a ponta.

2.31.3.2. Quando o yield se esgota

Chegar ao fim de uma função geradora (ou atingir um return explícito) levanta StopIteration automaticamente. Não é necessário levantá-la manualmente; o laço for ao redor a percebe e encerra.

Use geradores quando o código produtor for naturalmente escrito como um laço com alguns pontos de yield; use uma compreensão de lista simples quando você realmente precisar de toda a sequência na memória.