Списковые включения в Python мощнее, чем можно подумать

В Python списковые включения (и генераторы списков) —  замечательные механизмы, способные серьёзно упрощать программный код. Правда, чаще всего их используют в форме, предусматривающей наличие единственного цикла for и, возможно, одного условия if. И это всё. Но если попытаться немного вникнуть в эту тему, то окажется, что у списковых включений Python имеется гораздо больше возможностей, чем можно подумать, возможностей, разобравшись с которыми, можно, по меньшей мере, кое-чему научиться.

Множественные условия

Известно, что для фильтрации результатов работы спискового включения можно пользоваться условным оператором if. Обычно, если речь идёт о простом списковом включении, для достижения некоей цели достаточно одного if. А что если нам нужно вложенное условие?

values = [True, False, True, None, True]
print(['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:

print([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]

Если взглянуть на развёрнутый код, приведённый выше, становится понятно, что необязательно писать его именно так. Но синтаксис Python это позволяет.

Одна из причин, по которой может быть решено использовать именно такой подход, связана с читабельностью кода:

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

Уход от повторяющихся вычислений

Предположим, имеется списковое включение, вызывающее «тяжёлую» функцию и при проверке условия, и в теле цикла:

def func(val):
    # Тяжёлые вычисления...
    return val > 4

values = [1, 4, 3, 5, 12, 9, 0]
print([func(x) for x in values if func(x)])  # Неэффективно
# [True, True, True]

Такой подход неэффективен, так как он ведёт к удвоению времени вычислений. Можно ли как-то это исправить? Решить эту проблему нам помогут вложенные списковые включения!

print([y for y in (func(x) for x in values) if y])  # Эффективно
# [True, True, True]

Мне хотелось бы особо выделить то, что вышеприведённый код — это не двойной цикл. В этом примере мы создаём внутри спискового включения генератор, который используется внешним циклом. Если вам тяжело это читать — альтернативой такому коду может стать так называемый «моржовый» оператор:

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

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

Обработка исключений

Несмотря на то, что списковые включения обычно используются для решения простых задач — таких, как обработка каждого элемента списка с помощью какой-нибудь функции — бывает так, что внутри спискового включения может быть вызвано исключение. Однако не существует стандартного способа обработки исключений в списковых включениях. Как быть?

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"]
print([catch(int, value) for value in values])
print([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), авторы которого попытались предложить синтаксические конструкции для решения этой задачи, было отклонено.

Досрочный выход из цикла

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

print([i for i in iter(iter(range(10)).__next__, 4)])
# [0, 1, 2, 3]

from itertools import takewhile
print([n for n in takewhile(lambda x: x != 4, range(10))])
# [0, 1, 2, 3]

В первом из вышеприведённых примеров используется малоизвестная особенность функции iter. Конструкция iter(callable, sentinel) возвращает итератор, который «прерывает» итерацию сразу после того, как значение, возвращаемое функцией callable, окажется равным значению-метке sentinel. Когда внутренний блок iter возвращает значение-метку (в данном примере — 4) цикл автоматически останавливается.

Такой код не отличается хорошей читабельностью. Поэтому для решения той же задачи можно воспользоваться замечательным модулем itertools и функцией takewhile. Этот подход показан во втором из вышеприведённых примеров.

Хочу заметить, что если вы думали, что остановка цикла в списковом включении была возможна и раньше, это значит, что вы были правы. До Python 3.5 можно было воспользоваться вспомогательной функцией для выдачи исключения StopIteration внутри спискового включения. Это, правда, изменилось с принятием PEP 479.

Хитрости (и хаки)

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

Хотя и простое, безыскусное списковое включение — это очень мощный инструмент, его можно сделать ещё мощнее, если увязать его с библиотеками, наподобие 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)]
# ]

Тут имеется список дат, некоторые из которых следуют друг за другом, образуя непрерывные последовательности. Мы передаём даты функции consecutive_groups, используя для упорядочивания порядковые значения дат. Затем мы собираем возвращённые группы в список, пользуясь списковым включением.

В Python очень легко реализуется подсчёт накопительных сумм чисел. Можно просто передать itertools.accumulate список, а на выходе получатся суммы. А что если надо отменить подобную операцию?

from itertools import accumulate

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

print([y - x for x, y in more_itertools.pairwise(cumulative)])
# [4, 5, 12, 8, 1, 10, 21]

С помощью more_itertools.pairwise решение этой задачи выглядит до крайности простым.

Как уже было сказано, довольно-таки новый «моржовый» оператор можно использовать со списковыми включениями для создания локальных переменных. Это может пригодиться во многих ситуациях. Одна из них — применение функций any() и all().

Функция any() проверяет, удовлетворяет ли хотя бы одно значение в некоем итерируемом объекте определённому условию. Функция all() проверяет, удовлетворяют ли условию все такие значения. А что если нужно ещё и захватить значение, из-за которого any() возвращает True (так называемое «свидетельство»), или значение, из-за которого all() терпит неудачу (так называемый «контрпример»)?

numbers = [1, 4, 6, 2, 12, 4, 15]

# Возвращает лишь логические, но не числовые значения
print(any(number > 10 for number in numbers))  # True
print(all(number < 10 for number in numbers))  # False

# ---------------------
any((value := number) > 10 for number in numbers)  # True
print(value)  # 12

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

И any(), и all() используют сокращённый порядок вычислений при обработке переданных им данных. Это значит, что они останавливают работу, как только находят, соответственно, первое «свидетельство», или первый «контрпример». В результате, благодаря этой хитрости, переменная, созданная «моржовым» оператором, всегда даст нам первое «свидетельство» или первый «контрпример».

Итоги

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

языка. Даже тогда, когда речь идёт о чём-то таком, что может оказаться не особенно полезным в повседневной работе. И, кроме прочего, в таких вещах просто интересно и увлекательно разбираться.

Учитывая это — надеюсь, что вы узнали сегодня что-то новое. И напоследок хочу кое о чём предупредить. Если вы решите использовать в своих списковых включениях нечто вроде сложных условных конструкций или досрочно прерываемых циклов — не удивляйтесь, если ваши сослуживцы начнут косо на вас поглядывать.

О, а приходите к нам работать? 🤗 💰

Мы в wunderfund.io занимаемся высокочастотной алготорговлей с 2014 года. Высокочастотная торговля — это непрерывное соревнование лучших программистов и математиков всего мира. Присоединившись к нам, вы станете частью этой увлекательной схватки.

Мы предлагаем интересные и сложные задачи по анализу данных и low latency разработке для увлеченных исследователей и программистов. Гибкий график и никакой бюрократии, решения быстро принимаются и воплощаются в жизнь.

Сейчас мы ищем плюсовиков, питонистов, дата-инженеров и мл-рисерчеров.

Присоединяйтесь к нашей команде.