Comprehensions

Listas, conjuntos e dicionários

A forma geral é [expressão for item in iterável if condição]. O mesmo vale para chaves {}, que criam conjuntos e dicionários:

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]

Quando a transformação tem uma escolha, a expressão condicional vai antes do for. O if do final só filtra:

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

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

Compreensões aninhadas

Dois for na mesma compreensão percorrem as duas dimensões, do externo para o interno, na mesma ordem em que você escreveria os laços:

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

Expressões geradoras

Trocando os colchetes por parênteses, você obtém uma expressão geradora, que não monta a lista na memória e produz um item de cada vez. Dentro de uma chamada como sum, os parênteses extras são opcionais:

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

Funções como sum, any, all, min e max consomem geradores, e any e all param assim que o resultado está decidido.

O limite da legibilidade

Minha regra: uma compreensão com mais de um for ou com mais de uma condição já merece virar um laço comum ou uma função com nome. Se eu preciso reler duas vezes para entender, o próximo leitor também precisará.