Python Geradores (yield)

Na aula Iteradores, a contagem regressiva precisou de uma classe inteira: __init__, __iter__, __next__, raise StopIteration. Funciona, mas é muita cerimônia para "entregar 3, depois 2, depois 1". O Python tem um atalho: a função geradora. É uma função comum com uma palavra nova, yield (em inglês, "entregar, render"), e o Python escreve o iterador para você.

A ideia: a cozinha que frita um pastel por vez

A padaria do bairro recebe uma encomenda de 100 pastéis. Há dois jeitos de trabalhar:

• Jeito da lista: fritar os 100, empilhar todos na bancada e só então entregar a pilha. Precisa de uma bancada enorme, e o cliente espera até o último ficar pronto.
• Jeito do gerador: fritar um, entregar, e esperar o cliente pedir o próximo. A bancada só precisa caber um pastel, e a cozinheira lembra onde parou ("já fiz 37").

Uma função comum com return trabalha do primeiro jeito: calcula tudo e devolve de uma vez. Uma função com yield trabalha do segundo: entrega um item, pausa ali mesmo, e só continua quando alguém pede o próximo com next() (ou um for).

1. O primeiro gerador

A mesma contagem regressiva da aula anterior, agora como função geradora:

def contagem_regressiva(inicio):
    while inicio > 0:
        yield inicio
        inicio -= 1


for numero in contagem_regressiva(3):
    print(numero)
print("Já!")

$ python3 contagem.py
3
2
1
Já!

def contagem_regressiva(inicio):

Parece uma função comum (aula Funções). O que a torna geradora é ter um yield em algum lugar do corpo.

while inicio > 0:

Repete enquanto houver número para entregar.

yield inicio

Entrega o número da vez para quem pediu e pausa a função nesta linha. As variáveis (inicio) ficam guardadas.

inicio -= 1

Só roda quando alguém pede o próximo: a função acorda logo depois do yield.

o fim da função

Quando o while acaba e a função chega ao fim, o Python lança o StopIteration sozinho. Você não escreve o raise.

São 4 linhas no lugar das 13 da classe. Compare:

Chamar a função não roda a função

Isto é o que mais surpreende no começo. Chamar contagem_regressiva(2) não executa nenhuma linha do corpo: só cria um objeto gerador, parado antes da primeira linha. Ele é um iterador como os da aula anterior, então funciona com next():

def contagem_regressiva(inicio):
    while inicio > 0:
        yield inicio
        inicio -= 1


g = contagem_regressiva(2)
print(type(g))
print(iter(g) is g)
print(next(g))
print(next(g))
print(next(g))

$ python3 na_mao.py
<class 'generator'>
True
2
1
Traceback (most recent call last):
  File "/home/ana/na_mao.py", line 12, in <module>
    print(next(g))
          ~~~~^^^
StopIteration

• <class 'generator'>: o que a chamada devolveu foi um gerador, e não um número.
• iter(g) is g é True: como todo iterador, o gerador é o próprio marcador de página.
• O terceiro next(g), na linha 12, encontrou a função terminada: StopIteration, igualzinho ao da classe. O ~~~~^^^ aponta a chamada.

2. Vendo a pausa acontecer

Para enxergar quando cada pedaço do corpo roda, ponha prints dentro do gerador. O painel de senhas da padaria:

def painel():
    print("  [painel: ligando]")
    yield 101
    print("  [painel: chamando a próxima]")
    yield 102
    print("  [painel: desligando]")


senhas = painel()
print("Criei o gerador; nada rodou ainda.")
print("Recebi", next(senhas))
print("Recebi", next(senhas))
for resto in senhas:
    print("Recebi", resto)
print("Fim.")

$ python3 pausa.py
Criei o gerador; nada rodou ainda.
  [painel: ligando]
Recebi 101
  [painel: chamando a próxima]
Recebi 102
  [painel: desligando]
Fim.

1. senhas = painel(): cria o gerador. O "[painel: ligando]" não aparece aqui.
2. Primeiro next(): a função roda do começo até o primeiro yield. Aparece [painel: ligando] e o 101 é entregue.
3. Segundo next(): a função acorda depois do yield 101, mostra [painel: chamando a próxima] e para no yield 102.
4. O for pede o próximo: a função acorda, mostra [painel: desligando] e chega ao fim sem outro yield. O StopIteration automático encerra o for sem nenhuma volta.

def pares(limite):
    numero = 0
    while numero <= limite:
        print(numero)
        numero += 2


print(list(pares(10)))

def pares(limite):
    numero = 0
    while numero <= limite:
        yield numero
        numero += 2


print(list(pares(10)))

No código inicial, a função não tem yield: é uma função comum, que mostra os números e devolve None. O list(None) dá TypeError: 'NoneType' object is not iterable. Com o yield, ela vira geradora e o list() recolhe [0, 2, 4, 6, 8, 10].

3. yield × return

As duas funções abaixo calculam os quadrados de 0 a 4. A primeira monta uma lista e devolve com return; a segunda entrega um por um com yield:

def quadrados_lista(n):
    resultado = []
    for i in range(n):
        resultado.append(i * i)
    return resultado


def quadrados_gerador(n):
    for i in range(n):
        yield i * i


print(quadrados_lista(5))
print(quadrados_gerador(5))
print(list(quadrados_gerador(5)))
print(sum(quadrados_gerador(5)))

$ python3 lista_vs_gerador.py
[0, 1, 4, 9, 16]
<generator object quadrados_gerador at 0x7f6928110ee0>      « o endereço muda a cada execução
[0, 1, 4, 9, 16]
30

O print do gerador não mostra os números, porque eles ainda não foram calculados: só existe a "cozinha" pronta para fritar. Para ver tudo, peça com list(); para somar, o sum() pede um por um, sem montar lista nenhuma.

4. Expressões geradoras: ( … for … )

Você já conhece a list comprehension, entre colchetes. Troque os colchetes por parênteses e ela vira uma expressão geradora: o mesmo cálculo, mas entregue um por vez, sem montar a lista.

precos = [4.50, 12.00, 7.25, 30.00]

em_lista = [p * 2 for p in precos]
em_gerador = (p * 2 for p in precos)

print(em_lista)
print(em_gerador)
print(list(em_gerador))

total = sum(p for p in precos)
caros = sum(1 for p in precos if p > 10)
print("Total:", total)
print("Produtos acima de R$ 10:", caros)

$ python3 expressao.py
[9.0, 24.0, 14.5, 60.0]
<generator object <genexpr> at 0x7ff4c1f10ee0>
[9.0, 24.0, 14.5, 60.0]
Total: 53.75
Produtos acima de R$ 10: 2

[p * 2 for p in precos]

Colchetes: lista pronta, com os quatro valores na memória.

(p * 2 for p in precos)

Parênteses: um gerador. O nome que o Python dá a ele é <genexpr> ("expressão geradora").

sum(p for p in precos)

Quando a expressão geradora é o único argumento de uma função, os parênteses da função bastam: não precisa escrever sum((p for p in precos)).

sum(1 for p in precos if p > 10)

Um truque comum para contar: entrega um 1 para cada preço que passa no if, e o sum soma os uns.

feira = [("banana", 6.90, 2), ("tomate", 8.50, 1.5), ("alface", 3.00, 3)]

subtotais = [preco * qtd for nome, preco, qtd in feira]
print(subtotais)

feira = [("banana", 6.90, 2), ("tomate", 8.50, 1.5), ("alface", 3.00, 3)]

total = sum(preco * qtd for nome, preco, qtd in feira)
print(f"Total da feira: {total:.2f}")

5. Quanta memória cada um gasta

A diferença entre lista e gerador aparece de verdade com muitos itens. A função sys.getsizeof(), do módulo sys, diz quantos bytes um objeto ocupa na memória:

import sys

em_lista = [n * n for n in range(1_000_000)]
em_gerador = (n * n for n in range(1_000_000))

print("Lista:   ", sys.getsizeof(em_lista), "bytes")
print("Gerador: ", sys.getsizeof(em_gerador), "bytes")
print(sum(em_lista) == sum(em_gerador))

$ python3 memoria.py
Lista:    8448728 bytes
Gerador:  208 bytes
True

• A lista de um milhão de quadrados ocupa uns 8,4 milhões de bytes (8 MB), e isso é só a "estante" com o lugar de cada item; os números em si ocupam mais um tanto.
• O gerador ocupa 208 bytes, e ocuparia o mesmo com um bilhão de itens: ele guarda só a receita e o ponto onde parou.
• A última linha confirma que os dois dão a mesma soma. (O 1_000_000 é só um milhão escrito com sublinhados para ficar fácil de ler; veja a aula Números.)

import sys

em_lista = [n * n for n in range(1_000_000)]
em_gerador = (n * n for n in range(1_000_000))

print("Lista:   ", sys.getsizeof(em_lista), "bytes")
print("Gerador: ", sys.getsizeof(em_gerador), "bytes")

Quando usar cada um? Se você vai percorrer os itens uma vez (somar, procurar o maior, gravar num arquivo), o gerador basta e economiza memória. Se vai precisar deles várias vezes, acessar pela posição (itens[3]) ou saber quantos são (len), use uma lista.

6. Geradores infinitos

Como o gerador só calcula quando alguém pede, ele pode até não ter fim. O número do próximo pedido da lanchonete, por exemplo, nunca acaba:

from itertools import islice


def numeros_de_pedido():
    numero = 1
    while True:
        yield numero
        numero += 1


pedidos = numeros_de_pedido()
print(next(pedidos))
print(next(pedidos))
print(list(islice(pedidos, 5)))
print(next(pedidos))

$ python3 infinito.py
1
2
[3, 4, 5, 6, 7]
8

while True: yield numero

Um laço sem fim. Numa função comum, ele travaria o programa. Aqui, cada volta pausa no yield e só continua quando alguém pede.

from itertools import islice

Traz a função islice ("fatia de iterador") do módulo itertools, assunto da aula O Módulo itertools.

islice(pedidos, 5)

"Só os próximos 5". É o jeito de fatiar um gerador, que não aceita pedidos[:5]. O list() recolhe os 5 numa lista.

o último next(pedidos)

Deu 8: o gerador lembra onde parou, mesmo depois do islice.

Outro jeito de parar é um for com break (aula break e continue):

def numeros_de_pedido():
    numero = 1
    while True:
        yield numero
        numero += 1


for pedido in numeros_de_pedido():
    if pedido > 3:
        break
    print("Pedido", pedido)

$ python3 infinito_for.py
Pedido 1
Pedido 2
Pedido 3

7. Linha de montagem: um gerador alimentando outro

Geradores se encaixam como as estações de uma linha de montagem: cada um pega o que o anterior entrega, faz a sua parte e passa adiante. A padaria anotou as vendas do dia num arquivo, uma por linha (produto; preço; quantidade):

pão;0.75;12
leite;5.49;3
café;18.90;1
manteiga;12.50;2

def ler_linhas(nome):
    with open(nome, encoding="utf-8") as arquivo:
        for linha in arquivo:
            yield linha.strip()


def valores(linhas):
    for linha in linhas:
        produto, preco, quantidade = linha.split(";")
        yield float(preco) * int(quantidade)


linhas = ler_linhas("vendas.txt")
totais = valores(linhas)
print(f"Vendas do dia: R$ {sum(totais):.2f}")

$ python3 linha_de_montagem.py
Vendas do dia: R$ 69.37

1. ler_linhas abre o arquivo e entrega uma linha por vez, já sem o \n.
2. valores recebe esse gerador, separa cada linha no ; (o split da aula Métodos de Texto) e entrega preço × quantidade.
3. sum(totais) puxa a corrente: pede um valor a valores, que pede uma linha a ler_linhas, que lê uma linha do arquivo. Em nenhum momento o arquivo inteiro fica na memória.

Com 4 linhas tanto faz; com um arquivo de vendas do ano inteiro, de muitos megabytes, o programa gasta a mesma memória pequena.

yield from: repassar outro gerador

Quando um gerador quer entregar tudo o que outro entrega, em vez de escrever for item in outro(): yield item, use yield from ("entregar a partir de"):

def frutas():
    yield "banana"
    yield "maçã"


def legumes():
    yield "cenoura"


def feira():
    yield from frutas()
    yield from legumes()
    yield "ovos"


print(list(feira()))

$ python3 yield_from.py
['banana', 'maçã', 'cenoura', 'ovos']

Formas diferentes de fazer a mesma sequência

Os quadrados de 0 a 4, de quatro jeitos:

Quando dá errado

def contagem_regressiva(inicio):
    while inicio > 0:
        yield inicio
        inicio -= 1


numeros = contagem_regressiva(3)
print(numeros[0])

$ python3 erro_indice.py
Traceback (most recent call last):
  File "/home/ana/erro_indice.py", line 8, in <module>
    print(numeros[0])
          ~~~~~~~^^^
TypeError: 'generator' object is not subscriptable

$ python3 erro_len.py
Traceback (most recent call last):
  File "/home/ana/erro_len.py", line 8, in <module>
    print(len(numeros))
          ~~~^^^^^^^^^
TypeError: object of type 'generator' has no len()

def contagem_regressiva(inicio):
    while inicio > 0:
        yield inicio
        inicio -= 1


for numero in contagem_regressiva:
    print(numero)

$ python3 erro_sem_chamar.py
Traceback (most recent call last):
  File "/home/ana/erro_sem_chamar.py", line 7, in <module>
    for numero in contagem_regressiva:
                  ^^^^^^^^^^^^^^^^^^^
TypeError: 'function' object is not iterable

def contagem_regressiva(inicio):
    while inicio > 0:
        yield inicio
        inicio -= 1


numeros = contagem_regressiva(3)
print("Maior:", max(numeros))
print("Soma:", sum(numeros))

$ python3 gasto.py
Maior: 3
Soma: 0

Referência rápida

Exercícios

Missão final: as senhas do atendimento

from itertools import islice


def senhas():
    numero = 1
    # transforme isto num gerador infinito: while True, yield e numero += 1
    return f"A{numero:03d}"


print(list(islice(senhas(), 3)))

from itertools import islice


def senhas():
    numero = 1
    while True:
        yield f"A{numero:03d}"
        numero += 1


print(list(islice(senhas(), 3)))

O código inicial devolve o texto "A001", e um texto também é iterável: o islice pegou as 3 primeiras letras. Com o gerador, ele pega as 3 primeiras senhas. O {numero:03d} é a formatação com três dígitos e zeros à esquerda (aula f-strings).

Resumo

• Uma função com yield é uma função geradora: chamá-la não roda o corpo, só cria um gerador, que é um iterador.
• Cada next() roda o corpo até o próximo yield, entrega o valor e pausa, guardando as variáveis. No fim da função, o StopIteration vem sozinho.
• (x for x in itens), entre parênteses, é uma expressão geradora; como único argumento, dispensa os parênteses extras: sum(x for x in itens).
• Um gerador ocupa quase nada de memória, não importa quantos itens entregue; mas não tem [i], não tem len e se gasta numa passada.
• Geradores podem ser infinitos (while True): limite com islice, next ou break, nunca com list.
• Geradores se encadeiam numa linha de montagem; yield from repassa tudo de outro gerador.