itertools en Python: los bucles que dejas de escribir

itertools viene en la biblioteca estándar, así que no hay nada que instalar. Todo lo que hay adentro toma iteradores y devuelve iteradores, lo que significa que se combina con los generadores del post anterior y se mantiene igual de perezoso.

islice: tomar un pedazo de algo interminable

No puedes rebanar un generador.

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

Imprime:

TypeError: 'generator' object is not subscriptable

islice es la rebanada que funciona pidiendo, en lugar de por índice.

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)))

Imprime:

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

Los argumentos se leen como los de range: un final solo, o un inicio y un final. No puede contar hacia atrás, porque no hay forma de retroceder dentro de un iterador.

chain: un solo bucle sobre varias cosas

from itertools import chain

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

Imprime:

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

Una lista, una tupla y un range, recorridos como una sola secuencia, sin construir antes una lista combinada.

Cuando lo que tienes es una lista de listas, chain.from_iterable la aplana un nivel:

from itertools import chain

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

Imprime:

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

pairwise: cada elemento y el anterior

La diferencia entre lecturas consecutivas es un bucle que casi todo el mundo escribe con un índice y un error de uno.

from itertools import pairwise

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

Imprime:

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

Cinco pares a partir de seis lecturas, que es el número que querías. pairwise necesita Python 3.10 o más nuevo.

groupby, y la regla que no te puedes saltar

groupby agrupa elementos consecutivos. No junta todo lo que comparte una clave a lo largo de la secuencia: empieza un grupo nuevo cada vez que la clave cambia.

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])

Imprime:

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

Dos roles entraron, cuatro grupos salieron. No falló nada: ese es el comportamiento documentado, y es la razón por la que la gente concluye que groupby está roto.

Ordena primero por la misma clave y hace lo que querías:

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])

Imprime:

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

La misma función va a sorted y a groupby. Si esas dos alguna vez discrepan, vuelves a obtener el resultado de cuatro grupos.

Hay una segunda trampa. Cada grupo es un iterador sobre la misma secuencia de fondo, y solo es válido hasta que pasas al grupo siguiente:

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))

Imprime:

eng []
math []

El list() llegó hasta el último grupo antes de que se leyera nada, y todos los grupos anteriores quedaron atrás. Consume cada grupo dentro del bucle, o arma un diccionario real sobre la marcha.

Si solo quieres conteos, Counter del post sobre collections es más corto que todo esto.

count, cycle, repeat

Tres interminables. Solo se pueden usar junto con algo que se detenga: islice, un break, o un zip contra una secuencia 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))))

Imprime:

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

cycle guarda una copia de todo lo que ya vio, así que es el único miembro de este grupo que sí crece.

combinations y product

Dos bucles anidados que ya no tienes que anidar.

from itertools import combinations, product

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

Imprime:

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

combinations da cada par sin orden una sola vez. product es toda combinación de bucles anidados, y repeat=2 significa dos bucles sobre la misma secuencia.

Todos devuelven iteradores

Cada ejemplo de arriba está envuelto en list() por una razón: sin eso obtienes un objeto, no valores.

from itertools import chain

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

Imprime:

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

La dirección va a ser distinta en tu computadora. La segunda línea vacía no. Una pasada, igual que todo lo demás en esta serie.

Qué recordar

• islice(it, n) es cómo tomas algo de una fuente infinita. La sintaxis de rebanado no funciona sobre un iterador.

• chain recorre varias secuencias como una sola; chain.from_iterable aplana una lista de listas un nivel.

• pairwise te da cada elemento junto con el anterior, y saca bien la cuenta.

• groupby agrupa rachas, no valores. Ordena primero por la misma clave, y consume cada grupo antes de seguir.

• Todo esto es perezoso y de una sola pasada. Envuélvelo en list() cuando de verdad quieras los valores.