🔁 Как писать «правильные» циклы на Python: разбираемся на примерах

Несколько соображений о том, как сделать код циклов Python более производительным и ясным с помощью функций встроенного модуля itertools: product, isslice, takewhile.

Синтаксис оператора цикла в Python одновременно и прост, и не традиционен. По сравнению с C-подобными языками циклы в Python лишены общей трехступенчатой структуры for (init, condition, increment). В большинстве случаев достаточно for <item> in <iterable>. Цикл while <condition> используется реже.

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

Неужели цикл может быть «неправильным»?

У всякого языка программирования есть более или менее удачные подходы решения одних и тех же задач. Представим, вы спросили кого-то, кто только изучает Python: «Как получить текущий индекс при обходе списка?». Ответ может быть следующим:

index = 0

for name in names: 
    print(index, name)
    index += 1

Хотя приведенный цикл верен, это решение не в стиле Python. Разработчик с трехлетним опытом предложит такой код:

for i, name in enumerate(names):
     print(i, name)

enumerate() – это встроенная функция Python, которая принимает итерируемый объект в качестве параметра, а затем возвращает новый объект – генератор кортежей вида (текущий индекс, текущий элемент). Это лучший способ для данного случая: используется более интуитивно понятный код, к тому же он и продуктивнее.

Копнем приведенный пример поглубже. Цикл for состоит из структуры for <item> in <iterable>. Левая половина присваивает значение переменной item. В правой половине находится итерируемый объект, в качестве которого мы использовали функцию enumerate(). Это подводит нас к первой рекомендации.

Рекомендация 1: используйте функции, изменяющие итерируемый объект

Использование декоратора для обработки итерируемых объектов может по-разному влиять на код цикла. Прекрасный пример – встроенный модуль itertools. Это набор инструментальных функций, содержащий множество полезных итерируемых объектов. О самом модуле мы писали в статье «Итерируем правильно». В этом материале мы рассмотрим примеры использования функций модуля в практических задачах.

Используйте product() для компактности

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

def find_twelve(num_list1, num_list2, num_list3):
    """Находит все комбинации чисел из трех списков,
    в сумме дающие 12"""

    for num1 in num_list1:
        for num2 in num_list2:
            for num3 in num_list3:
                if num1 + num2 + num3 == 12:
                    return num1, num2, num3

Чтобы оптимизировать такие циклы, выполняющие обход объектов, можно использовать функцию product(). Функция принимает несколько итерируемых объектов и создает их декартово произведение.

from itertools import product

def find_twelve2(num_list1, num_list2, num_list3):
    for n1, n2, n3 in product(lst1, lst2, lst3):
        if n1 + n2 + n3 == 12:
            return n1, n2, n3

По сравнению с предыдущим кодом, цикл, использующий product(), нуждается только в одном уровне вложенности. Код становится более лаконичным.

Используйте islice(), чтобы обрабатывать только часть объектов цикла

Рассмотрим файл с заголовками постов Reddit следующего вида:

py-guide: Python guidebook, written for humans.
---
Python 3 Death Clock
---
Run any Python Script with an Alexa Voice Command
---
<...>

Между каждой парой заголовков, присутствует разделитель ---, а нам нужны только заголовки. Основываясь на том, что мы уже знаем о функции enumerate(), можно отфильтровать разделители по нечетным номерам:

def parse_titles(filename):
   """Читаем имя статьи reddit из файла"""

   with open(filename,'r') as fp:
      for i, line in enumerate(fp):
          # Опускаем разделитель
          if i% 2 == 0;
             yield line.strip()

Однако использование функции islice() из библиотеки itertools позволяет изменить сам итерируемый объект и упростить код. Функция islice (seq, start, end, step) имеет почти те же параметры, что оператор среза (list[start:stop:step]). Установим значение параметра step в 2 (по умолчанию 1).

from itertools import islice

def parse_ttiles_v2(filename):
    with open(filename, 'r') as fp:
        # Устанавливаем step=2
        # упускаем разделитель '---'
        for line in islice(fp, 0, None, 2):
            yield line.strip()

Используйте takewhile вместо break

Иногда необходимо определить, надо ли закончить цикл в самом его начале. Например:

for user in users:
    # При появлении первого неквалифицированного пользователя
    # дальнейшая обработка не производится
    if not is_qualified(user):
        break
    # Выполняем обработку ... 

Для раннего прерывания циклов можно использовать функцию takewhile(). Функция takewhile (predicate, iterable) проходит по всем объектам из iterable и вызывает функцию предиката, передав текущий объект в качестве аргумента, и проверяет возвращаемый результат.

Если функция предиката возвращает True, генерируется объект и цикл продолжается, в обратном случае цикл прерывается.

from itertools import takewhile

for user in takewhile(is_qualified, users):
   # Выполняем обработку ...

В itertools есть и другие интересные функции, которые можно использовать вместе с циклами:

• функция chain() позволяет сделать плоскими двухуровневые вложенные циклы;
• функция zip_longest() может организовать цикл сразу по нескольким объектам.
  

Если вас интересуют другие функции, переходите за подробностями к официальной документации или к упоминавшейся публикации.

Используйте генераторы для написания своего декоратора

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

def sum_even_only(numbers):
   """Суммирует все четные числа"""

   result = 0
   for num in numbers:
       if num % 2 == 0:
           result += num
   return result

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

def even_only(numbers):
   for num in numbers:
      if num% 2 == 0:
         yield num

def sum_even_only_v2(numbers):
   """Суммирует все четные числа"""

   result = 0
   for num in even_only(numbers):
       result += num
   return result

После декорирования переменной numbers функцией even_only, функции sum_even_only_v2 не приходится фильтровать четные номера – остается только просуммировать.

Рекомендация 2: разберите сложные блоки кода на составляющие

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

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

import time
import datetime
 
def award_active_users_in_last_30days():
    """Получаем всех юзеров, которые вошли в систему с 8 вечера до 10 вечера 
    в выходные дни в течение последних 30 дней и отправляем им бонусные баллы
    """

    days = 30
    for days_delta in range(days):
        dt = datetime.date.today()-datetime.timedelta(days=days_delta)
        # 5: Saturday, 6: Sunday
        if dt.weekday() not in (5, 6):
            continue
 
        time_start = datetime.datetime(dt.year, dt.month, dt.day, 20, 0)
        time_end = datetime.datetime(dt.year, dt.month, dt.day, 23, 0)
 
        # Преобразование в unix timestamp, необходимую для последующих запросов ORM
        ts_start = time.mktime(time_start.timetuple())
        ts_end = time.mktime(time_end.timetuple())
 
        # Опрашиваем юзеров и отправляем 1000 бонусных баллов
        for record in LoginRecord.filter_by_range(ts_start, ts_end):
            # Здесь можно добавить сложную логику
            send_awarding_points(record.user_id, 1000)

Вышеупомянутая функция состоит из двух уровней. Ответственность внешнего цикла заключается в отслеживании времени посещения за последние 30 дней, и преобразовании его в формат UNIX timestamp. Эти две метки времени используются внутренним циклом для дальнейшей передачи.

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

Как справиться со сложными циклами

Каковы недостатки такого кода? В один прекрасный день выяснилось, что некоторые пользователи не спят после полуночи по выходным и сидят на сайте. Появилось новое требование: «отправить уведомление юзерам, которые вошли в систему между 3:00 и 5:00 в выходные дни за последние 30 дней».

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

Чтобы эффективно использовать код повторно, нужно отсоединить часть функции, отвечающей за «выбор времени», от тела цикла. И в этом нам поможет наш старый друг, функция-генератор.

Разделение тела цикла с помощью функции-генератора

Чтобы отвязать выбор времени от цикла, определим функцию-генератор gen_weekend_ts_ranges(), которая используется для генерации меток времени UNIX:

def gen_weekend_ts_ranges(days_ago, hour_start, hour_end):
    """Создаем временной диапазон суббота-воскресенье
    и возвращаем его в виде UNIX timestamp.
    """
    
    for days_delta in range(days_ago):
        dt = datetime.date.today()-datetime.timedelta(days=days_delta)
        # 5: Saturday, 6: Sunday
        if dt.weekday() not in (5, 6):
            continue
 
        time_start = datetime.datetime(dt.year, dt.month, dt.day, hour_start, 0)
        time_end = datetime.datetime(dt.year, dt.month, dt.day, hour_end, 0)
 
        # Преобразование в unix timestamp, необходимое для последующих запросов ORM
        ts_start = time.mktime(time_start.timetuple())
        ts_end = time.mktime(time_end.timetuple())
        yield ts_start, ts_end

С помощью новой функции-генератора старую задачу «разослать наградные баллы» и новую задачу «разослать уведомления» можно реализовать повторным использованием одного и того же цикла:

def award_active_users_in_last_30days_v2(): 
    """Отправляет бонусные баллы"""

    for ts_start, ts_end in gen_weekend_ts_ranges(30, hour_start=20, hour_end=23):
        for record in LoginRecord.filter_by_range(ts_start, ts_end):
            send_awarding_points(record.user_id, 1000)

Подводя итоги

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

Кратко о некоторых моментах:

• использование декораторов для изменения итерируемого объекта может улучшить код циклов;
• множество полезных функций, способных улучшить цикл содержит модуль itertools;
  
• используйте функции-генераторы для простого определения собственного декоратора;
  
• не забывайте разделять логику бизнес-задач при разрастании циклов;
• используйте функции-генераторы для разделения блоков кода в цикле, выполняющих разные задачи, чтобы повысить гибкость.
  

Библиотека программиста надеется, что найдете эти подходы такими же полезными, как и мы. Удачи в обучении!

Источники

• https://towardsdatascience.com/tips-for-writing-authentic-loops-in-python-80304311e07f