Comprehensions

Lists, sets, and dictionaries

The general form is [expressão for item in iterável if condição]. The same applies to braces {}, which create sets and dictionaries:

quadrados = [n ** 2 for n in range(6)]
pares = [n for n in range(10) if n % 2 == 0]
print(quadrados)
print(pares)

palavras = ["python", "go", "rust"]
tamanhos = {p: len(p) for p in palavras}
unicos = {len(p) for p in palavras}
print(tamanhos, sorted(unicos))

[0, 1, 4, 9, 16, 25]
[0, 2, 4, 6, 8]
{'python': 6, 'go': 2, 'rust': 4} [2, 4, 6]

When the transformation has a choice, the conditional expression goes before of the for. The if the end only filters:

rotulos = ["par" if n % 2 == 0 else "ímpar" for n in range(4)]
print(rotulos)

['par', 'ímpar', 'par', 'ímpar']

Nested comprehensions

Two for in the same comprehension traverse both dimensions, from the outer to the inner, in the same order you would write the loops:

matriz = [[1, 2, 3], [4, 5, 6]]
achatada = [n for linha in matriz for n in linha]
transposta = [[linha[i] for linha in matriz] for i in range(3)]
print(achatada)
print(transposta)

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

Generator expressions

By replacing the brackets with parentheses, you get a generator expression, which does not build the list in memory and produces one item at a time. Within a call like sum, the extra parentheses are optional:

soma = sum(n ** 2 for n in range(1000))
print(soma)

nomes = ["Ana", "Bia", "Caio"]
print(any(len(n) > 3 for n in nomes), all(n[0].isupper() for n in nomes))

332833500
True True

Functions like sum, any, all, min and max consume generators, and any and all stop as soon as the result is determined.

The limit of readability

My rule: a comprehension with more than one for or with more than one condition already deserves to become a common loop or a named function. If I need to reread it twice to understand, the next reader will also need to.