itertools em Python: os laços que você para de escrever

O itertools está na biblioteca padrão, então não tem nada para instalar. Tudo nele recebe iteradores e devolve iteradores, o que quer dizer que ele se combina com os generators do post anterior e continua igualmente preguiçoso.

islice: pegue um pedaço de algo sem fim

Você não pode fatiar um generator.

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

Ele imprime:

TypeError: 'generator' object is not subscriptable

O islice é a fatia que funciona pedindo, em vez de indexando.

from itertools import islice

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

print(list(islice(naturals(), 5)))
print(list(islice(naturals(), 10, 15)))

Ele imprime:

[1, 2, 3, 4, 5]
[11, 12, 13, 14, 15]

Os argumentos se leem como os do range: um stop sozinho, ou um start e um stop. Ele não consegue contar para trás, porque não existe como voltar dentro de um iterador.

chain: um laço só sobre várias coisas

from itertools import chain

a, b, c = [1, 2], (3, 4), range(5, 7)
print(list(chain(a, b, c)))

Ele imprime:

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

Uma lista, uma tupla e um range, percorridos como uma sequência só, sem montar antes uma lista combinada.

Quando o que você tem é uma lista de listas, o chain.from_iterable achata um nível:

from itertools import chain

rows = [['ada', 'grace'], ['alan'], ['edsger', 'barbara']]
print(list(chain.from_iterable(rows)))

Ele imprime:

['ada', 'grace', 'alan', 'edsger', 'barbara']

pairwise: cada item e o anterior

Diferença entre leituras consecutivas é um laço que quase todo mundo escreve com um índice e um erro de um.

from itertools import pairwise

temps = [12, 14, 13, 18, 18, 21]
for a, b in pairwise(temps):
    print(f"{a} -> {b}  {b - a:+d}")

Ele imprime:

12 -> 14  +2
14 -> 13  -1
13 -> 18  +5
18 -> 18  +0
18 -> 21  +3

Cinco pares a partir de seis leituras, que é o número que você queria. O pairwise precisa do Python 3.10 ou mais novo.

groupby, e a regra que você não pode pular

O groupby agrupa itens consecutivos. Ele não junta tudo que tem a mesma chave ao longo da sequência — ele começa um grupo novo toda vez que a chave muda.

from itertools import groupby

people = [('ada', 'eng'), ('alan', 'math'), ('grace', 'eng'), ('emmy', 'math')]

for role, group in groupby(people, key=lambda p: p[1]):
    print(role, [name for name, _ in group])

Ele imprime:

eng ['ada']
math ['alan']
eng ['grace']
math ['emmy']

Dois cargos entraram, quatro grupos saíram. Nada deu errado; esse é o comportamento documentado, e é o motivo de as pessoas concluírem que o groupby está quebrado.

Ordene pela mesma chave antes e ele faz o que você queria:

from itertools import groupby

people = [('ada', 'eng'), ('alan', 'math'), ('grace', 'eng'), ('emmy', 'math')]

def by_role(p):
    return p[1]

for role, group in groupby(sorted(people, key=by_role), key=by_role):
    print(role, [name for name, _ in group])

Ele imprime:

eng ['ada', 'grace']
math ['alan', 'emmy']

A mesma função vai para o sorted e para o groupby. Se essas duas discordarem, você recebe o resultado de quatro grupos de novo.

Tem uma segunda armadilha. Cada grupo é um iterador sobre a mesma sequência de baixo, e ele só é válido até você passar para o grupo seguinte:

from itertools import groupby

people = [('ada', 'eng'), ('grace', 'eng'), ('alan', 'math')]

groups = list(groupby(people, key=lambda p: p[1]))
for role, group in groups:
    print(role, list(group))

Ele imprime:

eng []
math []

O list() andou até o último grupo antes de qualquer coisa ser lida, e todos os grupos anteriores ficaram para trás. Consuma cada grupo dentro do laço, ou monte um dict de verdade conforme avança.

Se você só quer contagens, o Counter do post sobre collections é mais curto que tudo isso.

count, cycle, repeat

Três sem fim. Eles só são usáveis junto com algo que para — o islice, um break, ou um zip contra uma sequência finita.

from itertools import count, cycle, repeat, islice

print(list(islice(count(10, 5), 4)))
print(list(islice(cycle('ab'), 5)))
print(list(zip('abc', repeat(0))))

Ele imprime:

[10, 15, 20, 25]
['a', 'b', 'a', 'b', 'a']
[('a', 0), ('b', 0), ('c', 0)]

O cycle guarda uma cópia de tudo que já viu, então ele é o único membro desse grupo que cresce.

combinations e product

Dois laços aninhados que você não precisa mais aninhar.

from itertools import combinations, product

print(list(combinations('abc', 2)))
print(list(product([0, 1], repeat=2)))

Ele imprime:

[('a', 'b'), ('a', 'c'), ('b', 'c')]
[(0, 0), (0, 1), (1, 0), (1, 1)]

O combinations dá cada par sem ordem uma vez só. O product é toda combinação de laços aninhados, e repeat=2 quer dizer dois laços sobre a mesma sequência.

Todos devolvem iteradores

Todo exemplo acima está envolvido em list() por um motivo: sem isso você recebe um objeto, não valores.

from itertools import chain

c = chain([1, 2], [3])
print(c)
print(list(c))
print(list(c))

Ele imprime:

<itertools.chain object at 0x7f9038d17e80>
[1, 2, 3]
[]

O endereço vai ser diferente na sua máquina. A segunda linha vazia não vai — uma passada, igual a todo o resto desta série.

O que lembrar

• islice(it, n) é como você pega de algo infinito. A sintaxe de fatia não funciona num iterador.

• O chain percorre várias sequências como uma; o chain.from_iterable achata uma lista de listas em um nível.

• O pairwise te dá cada item junto com o anterior, e acerta a contagem.

• O groupby agrupa sequências repetidas, não valores. Ordene pela mesma chave antes, e consuma cada grupo antes de seguir.

• Tudo aqui é preguiçoso e de uma passada só. Envolva em list() quando você realmente quiser os valores.