List comprehensions более эффективны, чем кажется на первый взгляд

List comprehension (списковые включения) — отличный функционал, который может значительно упростить код. Обычно они используются вместо одного цикла for, возможно, с добавлением одного условия if. Однако, если погрузиться в них подробнее, то можно понять, что у list comprehensions есть много возможностей, о которых вы и не знали...

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

Очевидно, что можно использовать условие (if), чтобы "фильтровать" результаты list comprehension. В простых случаях одного условия обычно достаточно. Но что если вы хотите использовать вложенное условие?
В этом случае поможет "тернарный условный оператор". Однако это не самое элегантное решение, поэтому вам придется решить, стоит ли экономия нескольких строк кода читаемости.

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 в одном list comprehension:

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]

Кажется, что сокращенная запись не очень хорошо читаема, но мы можем это подправить, отформатировав код должным образом:

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

Избегайте повторных вычислений

Допустим, у вас есть comprehension, которое вызывает тяжелую функцию как в условном выражении, так и в теле цикла:

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, которая может быть использована в другой части выражения.

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

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

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), в котором пытались ввести соответствующий синтаксис, было отклонено.

Прерывание цикла

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

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) возвращает итератор, который "прерывает" итерацию, как только значение вызываемой функции равно значению sentinel. Когда внутренний iter возвращает sentinel (4 в примере), цикл автоматически останавливается.

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

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

Хитрости

В предыдущих разделах мы рассмотрели некоторые не очень распространённые особенности list comprehension, которые могут быть полезными (и не очень) в повседневной работе. А теперь рассмотрим пару хитростей, которые сделают ваш код ещё более эффективным.

Хотя обычные list comprehensions сами по себе весьма полезные, они становятся ещё лучше, если их совместить с библиотекой 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 это довольно просто!

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

Функции Python any() и all() могут проверять, удовлетворяют ли какие-либо или все значения в итерируемом объекте определённому условию. Что, если мы хотим также сохранить значение, вызвавшее возврат True функции any() (так называемый "свидетель") или значение, вызвавшее False в функции 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() делают ленивую проверку. Это означает, что они останавливают вычисление, как только находят первый "свидетель" или "контрпример" соответственно. Поэтому переменная, созданная моржовым оператором, всегда даст нам первый "свидетель" / "контрпример".

Заключение

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

Надеемся, что вы научились чему-то новому. Но имейте в виду, что если вы решите использовать составные условия или прерывание цикла внутри list comprehension, то ваши коллеги могут вас возненавидеть.

Источник: Martin Heinz