Comprensión de listas en Python

Una comprensión construye una lista a partir de un bucle en una sola expresión. No es solo más corta: también es más rápida que ir agregando dentro de un bucle, porque el append se hace en C.

Las partes van siempre en el mismo orden: la expresión, después el for, después cualquier if.

El bucle y la comprensión

Estos dos producen la misma lista.

# the loop and the comprehension do the same thing
squares = []
for n in range(6):
    squares.append(n * n)
print(squares)

print([n * n for n in range(6)])

Imprime:

[0, 1, 4, 9, 16, 25]
[0, 1, 4, 9, 16, 25]

Filtrar

Un if al final se queda solo con los elementos que quieres.

# filtering with if
print([n for n in range(20) if n % 3 == 0])

Imprime:

[0, 3, 6, 9, 12, 15, 18]

Elegir un valor

Un if/else es parte de la expresión, así que va adelante. Este hace tropezar a la gente porque las dos formas se parecen y van en lugares distintos.

# if/else goes before the for
print(['even' if n % 2 == 0 else 'odd' for n in range(5)])

Imprime:

['even', 'odd', 'even', 'odd', 'even']

Regla práctica: el if al final filtra, el if/else adelante elige.

Comprensiones de diccionario y de conjunto

Las llaves te dan un diccionario cuando escribes un par, y un conjunto cuando escribes un solo valor.

# dict and set comprehensions
words = ['apple', 'bread', 'apple', 'milk']
print({w: len(w) for w in words})
print({len(w) for w in words})

Imprime:

{'apple': 5, 'bread': 5, 'milk': 4}
{4, 5}

La versión con conjunto descartó la longitud repetida, como hace un conjunto.

Dos cláusulas for

Varias cláusulas for se leen de izquierda a derecha, en el mismo orden en que las anidarías en un bucle. Así se aplana una lista de listas.

# two for clauses read top to bottom
pairs = [(x, y) for x in 'ab' for y in (1, 2)]
print(pairs)

grid = [[1, 2], [3, 4], [5, 6]]
print([cell for row in grid for cell in row])

Imprime:

[('a', 1), ('a', 2), ('b', 1), ('b', 2)]
[1, 2, 3, 4, 5, 6]

Expresiones generadoras

Los paréntesis redondos te dan un generador en su lugar. Produce valores a demanda y no guarda la lista entera en memoria.

# a generator expression does not build the list
import sys

listcomp = [n * n for n in range(100000)]
genexp = (n * n for n in range(100000))
print('list bytes:', sys.getsizeof(listcomp))
print('generator bytes:', sys.getsizeof(genexp))
print('sum is the same:', sum(listcomp) == sum(genexp))

Imprime:

list bytes: 800984
generator bytes: 200
sum is the same: True

La misma respuesta, cuatro mil veces menos memoria. Usa un generador cuando le vayas a pasar el resultado directo a algo que lo consuma una sola vez, como sum, any o un bucle.

La variable del bucle se queda adentro

En Python 3 la comprensión tiene su propio ámbito, así que no va a sobrescribir una variable que ya tenías.

# the loop variable does not leak
n = 'untouched'
result = [n for n in range(3)]
print(result, n)

Imprime:

[0, 1, 2] untouched

Qué recordar

• El orden es: expresión, for, y después if.

• El if al final filtra. El if/else va antes del for.

• Las llaves construyen diccionarios y conjuntos, los paréntesis redondos construyen un generador.

• La variable del bucle no se escapa al ámbito que la rodea.

Si una comprensión necesita más de dos cláusulas o no entra en una línea, escribe el bucle. Las comprensiones dejan de convenir apenas hay que descifrarlas.