Сообщить об ошибке.

Расширенное использование генераторов списков в Python

Использование генераторов списков в Python может значительно упростить код. В большинстве случаев генераторы списков используются только для написания одного цикла for/in, и возможно, с добавлением одного условия для цикла, и все. Но если немного разобраться в этой замечательной конструкции, то можно обнаружить много других особенностей, которым можно научиться.

Содержание:


Несколько условий в генераторах списков.

Ни для кого не секрет, что для фильтрации результатов генератора списка можно использовать условие if/else, с простыми включениями обычно достаточно одного if.

Что, если требуется вложенное условие?

>>> values = [True, False, True, None, True]
>>> ['yes' if v is True else 'no' if v is False else 'unknown' for v in values]
# ['yes', 'no', 'yes', 'unknown', 'yes']

Можно создать вложенное условие, используя тернарный оператор, хотя это не совсем красивое решение с точки зрения понимания/анализа кода человеком.

Код выше эквивалентен:

result = []
for v in values:
    if v is True:
        result.append('yes')
    else:
        if v is False:
            result.append('no')
        else:
            result.append('unknown')

print(result)
# ['yes', 'no', 'yes', 'unknown', 'yes']

Помимо использования сложных условных выражений, в генераторе списка можно складывать несколько if.

Смотрим следующий код:

>>> [i for i in range(100) if i > 10 if i < 20 if i % 2]
# [11, 13, 15, 17, 19]

Этот код эквивалентен:

# расширенная запись
result = []
for i in range(100):
    if i > 10:
        if i < 20:
            if i % 2:
                result.append(i)

print(result)
# [11, 13, 15, 17, 19]

Глядя на компактную запись генератора списка, в голову приходит мысль "- выглядит запутанно...", хотя синтаксис это позволяет. Одна из причин, по которой не хочется писать длинные генераторы списков - это удобство чтения кода (генераторы списков выглядят запутанно и сложны для анализа кода), но этот момент можно обойти. Принимая во внимание PEP-8, в котором сказано, что строки обрамленные любыми скобками можно переносить, то генератор списков можно записать следующим образом:

[
  i for i in range(100)
       if i > 10
       if i < 20
       if i % 2
]

Обработка исключений в генераторе списка.

Несмотря на то, что генератор списка обычно используется для простых задач, таких как вызов функции для каждого элемента списка, существуют ситуации, когда внутри генератора могут возникать исключения. Собственного способа обработки исключения внутри генератора списка НЕТ. Что можно с этим сделать?

Нужна функция-обработчик, чтобы поймать исключение внутри генератора списка.

def catch(f, *args, handle=lambda e: e, **kwargs):
    try:
        return f(*args, **kwargs)
    except Exception as e:
        return handle(e)

>>> values = [1, "text", 2, 5, 1, "also-text"]
>>> [catch(int, value) for value in values]
# альтернативный синтаксис
>>> [catch(lambda: int(value)) for value in values]
# [
#   1,
#   ValueError("invalid literal for int() with base 10: 'text'"),
#   2,
#   5,
#   1,
#   ValueError("invalid literal for int() with base 10: 'also-text'")
# ]

Здесь создается функцию catch(), которая принимает функцию и ее аргументы. Если внутри catch генерируется исключение, оно возвращается. Это конечно не идеальное решение, учитывая, что нужна вспомогательная функция, но это лучшее, что можно сделать (решение взято из PEP 463, который был отклонен).

Досрочный выход из цикла генератора списка.

Еще одним ограничением генератора списка является невозможность разорвать цикл (т.е. преждевременный выход из цикла). Хотя изначально это невозможно, можно реализовать небольшой хак, который решит проблему:

# первый способ
>>> [i for i in iter(iter(range(10)).__next__, 4)]
# [0, 1, 2, 3]

# второй способ
>>> from itertools import takewhile
>>> [n for n in takewhile(lambda x: x != 4, range(10))]
# [0, 1, 2, 3]

В первом способе выше используется малоизвестное поведение функции iter(). Функция Iter(callable, sentinel) возвращает итератор, который 'разрывает' итерацию, как только значение вызываемой функции становится равным значению sentinel. Когда внутренняя функция iter() возвращает sentinel (в примере 4), цикл автоматически останавливается. Это не очень читабельно.

Во втором способе для прерывания цикла генератора списка используется отличный встроенный модуль itertools и его функция itertools.takewhile.

Дорогие операции и генератор списка.

Допустим, есть генератор списка, который вызывает дорогостоящую функцию два раза (в условном выражении и в теле цикла). Это может выглядеть следующим образом:

def func(val):
    # дорогостоящие вычисления...
    return val > 4

>>> values = [1, 4, 3, 5, 12, 9, 0]
# эта запись будет работать не эффективно
>>> [func(x) for x in values if func(x)]
# [True, True, True]

Запись генератора списка выше неэффективна, так как удваивает время вычислений, но что можно с этим поделать? Вложенные генераторы списков спешат на помощь!

# эффективно, `func(x)` вызывается 
# только один раз
[y for y in (func(x) for x in values) if y]
# [True, True, True]

Здесь необходимо подчеркнуть, что это не двойной цикл. В этом примере внутри генератора списка создается выражение-генератор (обратите внимание на круглые скобки), который используется внешним циклом. Если эту запись трудно читать, альтернативой может быть использование оператора walrus.

Смотрим код:

>>> [y for x in values if (y := func(x))]

Здесь func() вызывается только один раз, создавая локальную переменную y, которую можно использовать в других частях выражения.

Маленькие хитрости для генераторов списков.

Простые, ванильные списки очень эффективны, но они становятся еще лучше в сочетании с такими библиотеками, как itertools или его расширением - сторонним модулем more-itertools.

Допустим, нужно найти серии последовательных чисел, дат, букв, логических значений или любых других упорядоченных объектов. Можно элегантно решить эту проблему, соединив more_itertools.consecutive_groups() из more-itertools с генератором списка.

Вот код:

import datetime
# pip install more-itertools
import more_itertools

dates = [
    datetime.datetime(2020, 1, 15),
    datetime.datetime(2020, 1, 16),
    datetime.datetime(2020, 1, 17),
    datetime.datetime(2020, 2, 1),
    datetime.datetime(2020, 2, 2),
    datetime.datetime(2020, 2, 4)
]

groups = [list(group) for group in more_itertools.consecutive_groups(dates, ordering=lambda d: d.toordinal())]
# [
# [datetime.datetime(2020, 1, 15, 0, 0), datetime.datetime(2020, 1, 16, 0, 0), datetime.datetime(2020, 1, 17, 0, 0)],
# [datetime.datetime(2020, 2, 1, 0, 0), datetime.datetime(2020, 2, 2, 0, 0)],
# [datetime.datetime(2020, 2, 4, 0, 0)]
# ]

Здесь есть список дат dates, некоторые из которых идут подряд. Передаем даты в функцию more_itertools.consecutive_groups(), используя порядковые значения дат d.toordinal() для упорядочения ordering. Затем собираем возвращенные группы в список, используя генератор списка.

Вычисление накопления сумм чисел в Python очень просто. Можно просто передать список в itertools.accumulate(), и на выходе получить суммы. Но что, если нужно отменить накопление? С помощью more_itertools.pairwise() это реально просто!

>>> from itertools import accumulate
>>> data = [4, 5, 12, 8, 1, 10, 21]
>>> cumulative = list(accumulate(data, initial=100))
>>> cumulative
# [100, 104, 109, 121, 129, 130, 140, 161]

# отменяем накопление
>>> [y - x for x, y in more_itertools.pairwise(cumulative)]
# [4, 5, 12, 8, 1, 10, 21]

Оператор walrus можно использовать со списками для создания локальной переменной. Это может быть полезно во многих ситуациях. Одна из таких ситуаций с функциями any() и all().

Функции Python any() и all() могут проверять, удовлетворяют ли некоторые или все значения в некотором итерируемом объекте определенным условиям. Что если нужно зафиксировать значение, из-за которого функция any() вернула True (так называемый "свидетель"), или значение, вызвавшее ошибку all() ("контрпример")?

>>> numbers = [1, 4, 6, 2, 12, 4, 15]
# Возвращает только логическое значение
>>> any(number > 10 for number in numbers)
# True  
all(number < 10 for number in numbers)
# False

# фиксируем значение
>>> any((value := number) > 10 for number in numbers) 
# True
>>> value  
# 12

>>> all((counter_example := number) < 10 for number in numbers)  
# False
>>> counter_example
# 12

Как any(), так и all() используют замыкания для вычисления заданного выражения. Это означает, что они прекращают оценку, как только находят первого "свидетеля". Следовательно переменная, созданная "оператором моржа" walrus, всегда будет давать требуемый результат.