Função geradora com yield

yield: devolver sem terminar

Na aula anterior você viu que o for precisa de um iterador que responda a next() e levante StopIteration no fim. Escrever isso na mão dá trabalho. O yield é o atalho oficial do Python para criar iteradores sem esse esforço. Basta usar yield dentro de uma função e ela deixa de ser comum: vira uma função geradora. A diferença de mentalidade é grande. Uma função normal calcula tudo e devolve um resultado pronto com return. Uma função geradora produz itens aos poucos, um a cada yield, pausando entre eles.

Pense no yield como um botão de pausa que também entrega algo pela janela. Quando a execução chega a um yield, três coisas acontecem: o valor à direita do yield é entregue a quem está percorrendo, a função congela naquela linha exata, com todas as suas variáveis intactas, e o controle volta para quem pediu o item. No próximo next(), a função descongela e segue da linha seguinte ao yield, como se o tempo não tivesse passado. Esse congelamento e descongelamento é o que dá aos geradores o seu superpoder.

def contagem_regressiva(n):
    while n > 0:
        yield n        # entrega n e pausa aqui
        n = n - 1      # continua daqui no proximo pedido
    print("Fim!")

for numero in contagem_regressiva(3):
    print(numero)

# 3
# 2
# 1
# Fim!

A cada volta do while, yield entrega um número e pausa; o for pede o próximo.

yield e return lado a lado

O contraste com return deixa tudo mais nítido. Uma função com return roda até encontrar o return, devolve aquele valor e acaba: qualquer código depois do return não roda naquela chamada. Uma função com yield pode entregar muitos valores, um por vez, e continua viva entre as entregas. Um return dentro de uma função geradora não devolve valor para o for; ele apenas encerra a geração, equivalendo a um StopIteration. É a diferença entre entregar um pacote fechado e abrir uma torneira que pinga sob demanda.

def com_return():
    return [1, 2, 3]   # calcula tudo e devolve a lista pronta

def com_yield():
    yield 1            # entrega 1 e pausa
    yield 2            # entrega 2 e pausa
    yield 3            # entrega 3 e pausa

print(com_return())        # [1, 2, 3]
print(com_yield())         # <generator object com_yield at 0x...>
print(list(com_yield()))   # [1, 2, 3]

com_return devolve a lista na hora; com_yield devolve um gerador que produz sob demanda.

Repare no detalhe do exemplo: chamar com_yield() não imprime 1, 2 e 3. Ele devolve um objeto gerador, ainda parado, que só produz os valores quando alguém os pede. Foi por isso que envolvemos com list(), que percorre o gerador do início ao fim e junta tudo numa lista. No dia a dia você raramente faz isso, porque o mais comum é percorrer o gerador direto com um for, aproveitando exatamente que ele produz aos poucos. Guardar essa vantagem, produzir sob demanda em vez de tudo de uma vez, é o que a próxima aula explora a fundo.

Conferindo o entendimento

O ponto para levar desta aula é que yield muda o contrato da função. Ela para de ser algo que roda e devolve um resultado, e passa a ser uma fábrica que entrega itens quando pedidos, pausando entre eles. Chamar a função só constrói a fábrica, o gerador; nada roda até você começar a puxar os itens com um for ou com next(). Esse detalhe confunde no começo e vale fixar bem, porque é a base de todo o resto do módulo, incluindo as generator expressions e os pipelines de geradores encadeados.