Generators e yield em Python

O post anterior terminou com vinte linhas de classe para contar de três até um. Aqui está a mesma coisa.

Três linhas

def countdown(n):
    while n > 0:
        yield n
        n -= 1

print(list(countdown(3)))

for i in countdown(3):
    print(i)

Ele imprime:

[3, 2, 1]
3
2
1

Qualquer função com yield dentro é uma generator function. Chamar ela não roda o corpo — ela devolve um generator, e um generator é um iterador. O Python escreve o __iter__ e o __next__ por você.

Nada roda antes de você pedir

Essa é a parte que vale olhar com atenção.

def noisy():
    print('  starting')
    yield 1
    print('  between')
    yield 2
    print('  finishing')

g = noisy()
print('created — the body has not run')
print('got', next(g))
print('got', next(g))
try:
    next(g)
except StopIteration:
    print('done')

Ele imprime:

created — the body has not run
  starting
got 1
  between
got 2
  finishing
done

Repare na ordem. O starting aparece depois do created, não antes. O corpo roda só quando alguma coisa chama next, e ele para de novo no instante em que encontra um yield.

O return encerra uma função e joga fora tudo que era local a ela. O yield suspende a função com as variáveis intactas, e o next() seguinte continua na linha de baixo. Quando o corpo enfim termina, o Python levanta StopIteration — o mesmo sinal que o post anterior capturou na mão.

Por que ser preguiçoso vale alguma coisa

import sys

squares_list = [n * n for n in range(1_000_000)]
squares_gen  = (n * n for n in range(1_000_000))

print('list:', sys.getsizeof(squares_list), 'bytes')
print('gen: ', sys.getsizeof(squares_gen), 'bytes')
print('same total:', sum(squares_list) == sum(squares_gen))

Ele imprime:

list: 8448728 bytes
gen:  200 bytes
same total: True

A lista guarda um milhão de números. O generator guarda uma função pausada. Os dois somam o mesmo valor, e só um deles precisou caber na memória.

Parênteses no lugar de colchetes é a diferença inteira. Isso é uma generator expression — uma list comprehension que nunca monta a lista.

nums = [1, 2, 3, 4, 5, 6]

print(sum(n * n for n in nums if n % 2))
print(sorted((n * n for n in nums), reverse=True)[:3])

Ele imprime:

35
[36, 25, 16]

Quando a generator expression é o único argumento, você pode largar os parênteses dela — é por isso que sum(n * n for n in nums) se lê do jeito que se lê.

Agora infinito é permitido

Uma lista com todos os números naturais não é uma coisa que você pode ter. Um generator deles é tranquilo, porque ele só guarda aquele em que está.

def naturals():
    n = 1
    while True:
        yield n
        n += 1

out = []
for n in naturals():
    if n * n > 500:
        break
    out.append(n * n)

print(out)

Ele imprime:

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484]

O while True nunca termina. Quem termina é o break, e o generator fica simplesmente pausado, para sempre, e é coletado.

Estágios que puxam um do outro

Generators se encadeiam. Cada estágio pede um único item ao estágio anterior.

lines = ['17,ada,engineer', '', '18,grace,admiral', 'not a row', '19,alan,logician']

good   = (l for l in lines if l.count(',') == 2)
fields = (l.split(',') for l in good)
people = ((int(i), name) for i, name, _ in fields)

for pid, name in people:
    print(pid, name)

Ele imprime:

17 ada
18 grace
19 alan

Nada aconteceu até o laço for pedir. Aí uma linha passou pelos três estágios, foi impressa, e a próxima veio atrás. Em nenhum momento existiu uma lista intermediária inteira. Troque lines por um arquivo de 4GB e o código não muda.

yield from

Quando o trabalho de um generator é repassar o que outro produz, o yield from diz isso em uma linha.

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

print(list(flatten([1, [2, [3, 4]], 5, [[6]]])))

Ele imprime:

[1, 2, 3, 4, 5, 6]

Sem ele você escreveria um laço que re-emite cada item, o que funciona e se lê pior.

As duas coisas que você abre mão

Um generator é um iterador, então ele roda uma vez — tudo que o post anterior disse vale aqui.

g = (n for n in range(4))
print(list(g))
print(list(g))

Ele imprime:

[0, 1, 2, 3]
[]

E ele não tem tamanho e não tem indexação, porque ele não sabe o que ainda não produziu.

g = (n for n in range(5))
try:
    print(len(g))
except TypeError as err:
    print('TypeError:', err)

try:
    print(g[0])
except TypeError as err:
    print('TypeError:', err)

Ele imprime:

TypeError: object of type 'generator' has no len()
TypeError: 'generator' object is not subscriptable

Se você precisa do tamanho ou de um índice, você precisava de uma lista. Chame list() nele e pague a memória.

O que lembrar

• Um yield no corpo da função faz dela uma generator function. Chamar ela não roda nada.

• Cada next() roda o corpo até o próximo yield e congela ele ali, variáveis locais e tudo.

• Parênteses te dão uma generator expression: a comprehension que você já escreve, sem a lista.

• Generators podem ser infinitos, e eles se encadeiam em pipelines que nunca seguram mais de um item por estágio.

• Uma passada, sem len, sem indexação. Esse é o preço, e quase sempre vale pagar.