List Comprehensions em Python

Uma comprehension monta uma lista a partir de um laço, em uma única expressão. Ela não é só mais curta: também é mais rápida que dar append dentro de um laço, porque o append acontece em C.

As partes vêm sempre na mesma ordem: a expressão, depois o for, depois qualquer if.

O laço e a comprehension

Estes dois produzem a mesma 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)])

Ele imprime:

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

Filtrando

Um if no fim mantém só os itens que você quer.

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

Ele imprime:

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

Escolhendo um valor

Um if/else faz parte da expressão, então ele vai na frente. Esse é o que confunde as pessoas, porque as duas formas se parecem e ficam em lugares diferentes.

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

Ele imprime:

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

Regra prática: if no fim filtra, if/else na frente escolhe.

Comprehensions de dict e de set

Chaves te dão um dicionário quando você escreve um par, e um set quando você escreve um valor só.

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

Ele imprime:

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

A versão com set descartou o comprimento repetido, como um set faz.

Dois for na mesma comprehension

Vários for são lidos da esquerda para a direita, na mesma ordem em que você os aninharia num laço. É assim que você achata uma 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])

Ele imprime:

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

Generator expressions

Parênteses te dão um generator no lugar da lista. Ele produz os valores sob demanda e não guarda a lista inteira na memória.

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

Ele imprime:

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

Mesma resposta, quatro mil vezes menos memória. Use um generator quando você vai passar o resultado direto para algo que o consome uma vez só, como sum, any ou um laço.

A variável do laço fica lá dentro

No Python 3 a comprehension tem escopo próprio, então ela não sobrescreve uma variável que você já tinha.

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

Ele imprime:

[0, 1, 2] untouched

O que lembrar

• A ordem é: expressão, for, depois if.

• if no fim filtra. if/else vai antes do for.

• Chaves constroem dicionários e sets, parênteses constroem um generator.

• A variável do laço não vaza para o escopo em volta.

Se uma comprehension precisa de mais de duas cláusulas ou não cabe numa linha, escreva o laço. Comprehensions param de valer a pena no momento em que precisam ser decifradas.