Comprehensions: list, dict, set и generator

В первом случае условие фильтрует, во втором - выбирает что добавить.

Вложенные циклы

# Декартово произведение
pairs = [(x, y) for x in [1, 2, 3] for y in ["a", "b"]]
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b'), (3, 'a'), (3, 'b')]

# Развёртывание двумерного списка
matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flat = [num for row in matrix in for num in row]
# Опечатка выше для красоты - правильно:
flat = [num for row in matrix for num in row]
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

Порядок for слева направо - как во вложенных циклах. Это иногда удивляет: «снаружи» циклы, «внутри» выражение.

Dict comprehension

# Из списка кортежей
data = [("Alice", 30), ("Bob", 25), ("Charlie", 35)]
ages = {name: age for name, age in data}
# {'Alice': 30, 'Bob': 25, 'Charlie': 35}

# Из range
squares = {x: x ** 2 for x in range(5)}
# {0: 0, 1: 1, 2: 4, 3: 9, 4: 16}

# С фильтром
adults = {name: age for name, age in data if age >= 18}

Синтаксис: {ключ: значение for ... in ...}. Двоеточие между ключом и значением - то же что в литерале dict.

Set comprehension

unique_lengths = {len(word) for word in ["hello", "world", "hi", "bye"]}
# {5, 2, 3}

Внешне как list, но в фигурных скобках. Дедупликация бесплатно: одинаковые элементы сольются - это свойство множеств.

Generator expression - ленивый comprehension

Если результат не нужно материализовать в список, используй generator expression - такой же синтаксис, но в круглых скобках:

# List comprehension - вся коллекция в памяти
squares_list = [x ** 2 for x in range(1_000_000)]  # ~8 МБ

# Generator expression - ленивая, по запросу
squares_gen = (x ** 2 for x in range(1_000_000))   # ~200 байт

Generator не хранит результаты - вычисляет на лету. Полноценные генераторы с yield разберём отдельно. Идеально для:

• Однократного перебора больших данных
• Передачи в функции типа sum, min, max, any, all
• Pipelining

total = sum(x ** 2 for x in range(1_000_000))  # без промежуточного списка

sum(x ** 2 for x in range(10))   # OK
max(len(s) for s in strings)     # OK

some_func((x for x in items), other_arg)

Когда НЕ использовать comprehension

Comprehensions - сжатый код. Это не всегда хорошо. Не используй их:

1. Когда тело цикла длинное.

# Плохо
results = [some_complex_calculation(x, y, z) if x > 0 and y < 100 and z != None else default_value for x, y, z in data if x is not None and validation_passed(x)]

# Лучше - обычный цикл
results = []
for x, y, z in data:
    if x is None or not validation_passed(x):
        continue
    if x > 0 and y < 100 and z is not None:
        results.append(some_complex_calculation(x, y, z))
    else:
        results.append(default_value)

2. Когда есть side effects.

# Плохо
[send_email(user) for user in users]   # comprehension для side effect

# Хорошо
for user in users:
    send_email(user)

Comprehension - для создания коллекции. Если результат не нужен - используй обычный цикл.

3. Когда условий и циклов больше двух.

# Плохо - три уровня
[(x, y, z) for x in xs for y in ys if y > x for z in zs if z > y]

# Лучше - вложенные циклы
result = []
for x in xs:
    for y in ys:
        if y > x:
            for z in zs:
                if z > y:
                    result.append((x, y, z))

Множественная распаковка

Comprehension работает с распаковкой:

points = [(1, 2), (3, 4), (5, 6)]
xs = [x for x, _ in points]
sums = [x + y for x, y in points]

С enumerate:

indexed = {i: word for i, word in enumerate(["a", "b", "c"])}
# {0: 'a', 1: 'b', 2: 'c'}

С zip:

mapping = {key: value for key, value in zip(keys, values)}
# или короче через dict()
mapping = dict(zip(keys, values))

walrus в comprehension

С Python 3.8 можно использовать walrus :=:

# Без walrus - вычисляем ext дважды
files = [(path, get_ext(path)) for path in paths if get_ext(path) == "py"]

# С walrus - один раз
files = [(path, ext) for path in paths if (ext := get_ext(path)) == "py"]

Иногда удобно, но при злоупотреблении читается тяжело.

Производительность

Comprehensions быстрее ручных циклов с append - это и читаемо, и оптимально.

# Comprehension
results = [str(x) for x in numbers]

# map
results = list(map(str, numbers))

# Уникальные значения с сохранением порядка
seen = set()
unique = [x for x in items if not (x in seen or seen.add(x))]

# Альтернатива (Python 3.7+, dict сохраняет порядок)
unique = list(dict.fromkeys(items))

# Уникальные без порядка
unique = list({x for x in items})

[print(x) for x in items]   # анти-паттерн

items = [1, 2, 3]
items = [items.append(x) for x in items]   # мусор, не делай

# Превратить в список квадратов чётных чисел до 20
squares = []
for x in range(20):
    if x % 2 == 0:
        squares.append(x ** 2)

# Comprehension:
squares = [x ** 2 for x in range(20) if x % 2 == 0]

names = ["Alice", "Bob", "Charlie"]
ages = [30, 25, 35]

# Через dict comprehension
people = {name: age for name, age in zip(names, ages)}

# Через dict() конструктор
people = dict(zip(names, ages))

# Сумма квадратов от 0 до 1_000_000
# Не создавай промежуточный список - используй generator expression
total = sum(x ** 2 for x in range(1_000_000))
print(total)