Comprensiones — un bucle en una línea

El problema que vamos a resolver

Una forma de código ha ido apareciendo una y otra vez en los últimos seis capítulos:

marks = [72, 45, 90, 33]

doubled = []
for mark in marks:
    doubled.append(mark * 2)
print(doubled)

doubled = [mark * 2 for mark in marks]
print(doubled)

[144, 90, 180, 66]
[144, 90, 180, 66]

Las cuatro líneas de arriba y la de abajo hacen exactamente lo mismo.

Tres de esas cuatro líneas no son el trabajo, son el papeleo — crear una lista vacía, recorrer el origen, hacer append cada vez. La única decisión de verdad es mark * 2. Todo lo demás es igual cada vez.

Una comprensión elimina esa repetición. No es una habilidad nueva — es lo que ya sabías escribir, con menos palabras. De ahí viene también su límite, y de eso trata la última sección de este capítulo: acorta lo que se vuelve más claro por estar acortado, no todo.

Al terminar este capítulo podrás

• Escribir y leer una comprensión de listas
• Filtrar con if y elegir con if/else
• Escribir comprensiones de diccionarios y de conjuntos
• Decir cuándo una comprensión es la opción equivocada
• Desmontar una comprensión ilegible de vuelta a un bucle

Requisitos previos: Datos anidados — una lista de diccionarios.

La forma

[ result   for  item   in  source ]

Aquí result es lo que quieres construir, item es una cosa cada vuelta, y source es la lista (u otra colección) de la que salen las cosas.

Leída en voz alta es una frase normal: «`mark 2 para cada mark en marks`»*.

Una forma de retenerlo: primero lee la parte del for, después el principio. Saber de dónde salen las cosas hace que lo que se está construyendo se vea fácil.

Filtrado

Un if al final conserva solo las cosas que cumplen la condición:

marks = [72, 45, 90, 33]

passed = [mark for mark in marks if mark >= 40]
print(passed)

[72, 45, 90]

Aquí la primera parte es solo mark — sin cálculo, tal cual. Filtrar y transformar se pueden hacer a la vez:

names = ["  rafi ", "AHMED", " dia"]

clean = [name.strip().lower() for name in names if name.strip()]
print(clean)

['rafi', 'ahmed', 'dia']

La condición es name.strip() sin ninguna comparación. Recuerda la regla del capítulo ocho: una cadena vacía es falsa. Así que alguien que solo escribió espacios se descarta.

if al final versus if al principio

Dos cosas distintas, y la posición es la diferencia:

marks = [72, 33, 90]

labels = ["pass" if m >= 40 else "fail" for m in marks]
print(labels)

['pass', 'fail', 'pass']

Un if al final filtra — algunas cosas nunca llegan al resultado, así que el resultado puede ser más corto.

Un if ... else al principio elige — cada cosa produce algo, así que la longitud no cambia. Es la condición en una línea del capítulo quince, dentro de una comprensión.

¿Cambian la longitud? Esa pregunta es la forma más fácil de distinguir las dos.

Diccionarios y conjuntos también

La misma forma, corchetes distintos:

marks = {"rafi": 72, "ahmed": 45, "bilal": 90}

passed = {name: mark for name, mark in marks.items() if mark >= 50}
print(passed)

doubled = {name: mark * 2 for name, mark in marks.items()}
print(doubled)

{'rafi': 72, 'bilal': 90}
{'rafi': 144, 'ahmed': 90, 'bilal': 180}

Llaves y dos puntos hacen una comprensión de diccionarios. marks.items() devuelve una tupla cada vuelta y name, mark la desempaqueta, igual que en un bucle.

Las mismas llaves sin dos puntos hacen un conjunto:

orders = [
    {"item": "pen"},
    {"item": "bag"},
    {"item": "pen"},
]

items = {order["item"] for order in orders}
print(sorted(items))

['bag', 'pen']

La pregunta del capítulo dieciséis de «cuántos distintos» en una línea. pen llega dos veces y aparece una.

De números a cadenas — algo que ya has visto

Esta línea apareció en el capítulo trece con la promesa de que la explicación estaba por venir:

marks = [72, 45, 90]
print(", ".join(str(m) for m in marks))
print([str(m) for m in marks])

72, 45, 90
['72', '45', '90']

Ahora se puede leer completa. str(m) for m in marks es una comprensión sin los corchetes cuadrados — a esa forma se la llama expresión generadora, y en lugar de construir la lista completa de una vez le entrega a join un elemento cada vez. Los corchetes se pueden omitir cuando es el único argumento de una función, y con datos grandes ahorra memoria.

Cuándo no escribir una

Saber escribir una comprensión no significa que debas hacerlo. Mira esto:

rows = [["1", "2"], ["3", "x"]]

numbers = [int(c) for row in rows for c in row if c.isdigit()]
print(numbers)

[1, 2, 3]

La línea es correcta. Pero tiene dos for y un if, y hay que recordar el orden: los for se leen de fuera hacia dentro, como se escribirían los bucles anidados. Seis meses después, esa línea cuesta trabajo leerla.

Como bucle:

rows = [["1", "2"], ["3", "x"]]

numbers = []
for row in rows:
    for c in row:
        if c.isdigit():
            numbers.append(int(c))

print(numbers)

[1, 2, 3]

Más líneas, menos pensamiento. El objetivo del código no es ser corto sino ser claro — y la mayoría de las veces acortarlo lo hace más claro, pero no siempre.

Como regla práctica: un for y como mucho un if — una comprensión. Más que eso — un bucle. Y nunca pongas un print ni un append dentro de una comprensión; es una cosa para construir un valor, no para realizar acciones.

Un ejemplo completo

views.py:

# One list of records in, three different views out
orders = [
    {"customer": "rafi", "item": "pen", "quantity": 3, "price": 15.0},
    {"customer": "ahmed", "item": "bag", "quantity": 1, "price": 850.0},
    {"customer": "rafi", "item": "ink", "quantity": 2, "price": 120.0},
    {"customer": "dia", "item": "pen", "quantity": 5, "price": 15.0},
]

totals = [order["quantity"] * order["price"] for order in orders]
print("Totals    :", totals)

big = [order["item"] for order in orders if order["quantity"] * order["price"] > 100]
print("Over 100  :", big)

items = {order["item"] for order in orders}
print("Distinct  :", sorted(items))

by_item = {order["item"]: order["price"] for order in orders}
print("Price list:", by_item)

print("Grand     :", sum(totals))

Totals    : [45.0, 850.0, 240.0, 75.0]
Over 100  : ['bag', 'ink']
Distinct  : ['bag', 'ink', 'pen']
Price list: {'pen': 15.0, 'bag': 850.0, 'ink': 120.0}
Grand     : 1210.0

Los mismos datos que en el capítulo diecisiete, con la mayor parte de lo que allí era un bucle ahora en una sola línea.

Tres cosas que vale la pena notar.

totals y sum(totals) son líneas separadas. El total podría haber sido sum(order["quantity"] * order["price"] for order in orders), pero entonces los valores individuales de cada línea no existirían. Un nombre intermedio suele valer más que un truco de una línea.

El mismo cálculo aparece dos veces dentro de big — una vez en la condición, y haría falta otra vez en el resultado. Aquí el resultado solo toma el nombre, así que no se pierde nada; pero cuando necesitas las dos cosas, esa es una buena razón para dejar la comprensión y escribir un bucle, donde el valor se puede calcular una vez y ponerle un nombre.

by_item perdió un registro. Cuatro pedidos, tres pares en el resultado — pen aparece dos veces, y por la regla del capítulo quince el último reemplazó en silencio al anterior. Aquí los dos pen tienen el mismo precio, así que no pasó nada; con precios distintos habría salido mal en silencio. Una comprensión no crea ese problema, pero al ser corta da menos ocasión de notarlo.

Cuando falla

SyntaxError por poner if y else al final [m for m in marks if m >= 40 else 0] no es válido. El if de filtrado va al final y no tiene else; el if ... else de elección va al principio: [m if m >= 40 else 0 for m in marks].

El resultado está lleno de None Algo dentro de la comprensión no devuelve nada — como [items.append(x) for x in other]. append devuelve None, así que el resultado es una lista de None. Una comprensión construye valores; para realizar acciones, escribe un bucle normal.

NameError — la variable de la comprensión ya no existe fuera Después de [m * 2 for m in marks] no puedes usar m. En el capítulo once la variable de un bucle for sobrevivía al bucle; la de una comprensión no — solo vive dentro.

Una comprensión de diccionarios produjo menos pares de los esperados La misma clave se produjo más de una vez y la última reemplazó a la anterior. Compara con len().

No puedo leer la línea Esa es razón suficiente. Rómpela en un bucle — la diferencia de velocidad casi nunca importa para quien está empezando, y la próxima persona en leerla, probablemente tú, lo agradecerá.

Tu turno

Escribe un archivo llamado clean.py con una lista de nombres desordenados — algunos con espacios de más, otros con mayúsculas, y al menos uno que esté totalmente en blanco.

Usando comprensiones, construye:

1. Una lista de nombres limpios — recortados, en minúsculas, sin los que están en blanco
2. Una lista con la longitud de cada nombre
3. Una lista con los nombres de más de cinco caracteres
4. Un diccionario con {name: length}
5. Un conjunto con las letras iniciales

Luego, dos experimentos:

• Escribe el número 1 de nuevo como un bucle for normal, y comprueba que los dos resultados son iguales con ==
• Escribe una comprensión con dos for y un if, y luego desmántala en un bucle. ¿Cuál se lee más fácil?

Esa última pregunta no tiene respuesta «correcta», y esa es la gracia. Aprender a hacerte esa pregunta es la lección de verdad de este capítulo.