2.31. Iteradores e geradores

O ciclo for tem feito mais trabalho do que parece. Esta página aborda o protocolo do iterador em que se baseia, e a palavra-chave yield que permite construir iteradores personalizados.

2.31.1. O protocolo do iterador

Todo o objeto que pode ser percorrido implementa dois métodos:

• __iter__() – devolve um iterador sobre os elementos do objeto.

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

O built-in 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 em __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 ...

Todas as listas, tuplos, strings, dicionários, conjuntos, objetos de ficheiro e geradores já implementam __iter__ e __next__ – razão pela qual 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; devolve um objeto gerador (um iterador) que executa o corpo um yield de cada 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é ao próximo yield, entrega esse valor ao chamador e pausa aí. O estado local (i neste caso) é preservado entre as retomas.

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

Os geradores são a forma mais simples de produzir uma sequência de forma lazy – não é construída nenhuma lista, os elementos são calculados apenas quando o consumidor os solicita, e a função pode produzir elementos indefinidamente se quiser.

2.31.3. Pipelines lazy

Os geradores compõem-se 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 intermédia, sem limite superior embutido em numbers, e o consumidor (for v in pipeline) decide quando parar.

Cada next() no consumidor desencadeia um pull através da cadeia; os valores existem apenas quando algo os solicita.

2.31.3.1. yield from

Um ciclo que obtém elementos de outro iterável e produz cada um é suficientemente comum para que o Python forneça um atalho. A expressão yield from iter produz cada valor que o iterável gera, por ordem – como se o gerador tivesse um ciclo for x in iter: yield x inline:

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 ciclo for explícito, apenas mais curto, e propaga StopIteration do iterável interno para o gerador externo de forma limpa – útil quando se encadeiam 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á-lo manualmente; o ciclo for circundante deteta-o e termina.

Use geradores quando o código de produção se escreve naturalmente como um ciclo com alguns pontos de yield; use uma compreensão de lista simples quando genuinamente precisa da sequência completa em memória.