🐍 Python enumerate: упрощаем циклы с помощью счетчиков

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

Перевод публикуется с сокращениями, автор оригинальной статьи Bryan Weber.

Итерация с помощью циклов for

Цикл for в Python использует итерацию на основе коллекции. Это означает, что Python на каждой итерации назначает следующий элемент из iterable переменной цикла, как в этом примере:

>>> values = ["a", "b", "c"]

>>> for value in values:
...     print(value)
...
a
b
c

Здесь значения состоят из трех строк: a, b и c. В Python списки – один из типов итерационных объектов iterable. В цикле for переменная имеет значение value. На каждой итерации ее значение устанавливается равным следующему элементу из values.

Преимущество такой итерации заключается в том, что она помогает избежать ошибки off-by-one, которая часто встречается в других ЯП.

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

>>> index = 0

>>> for value in values:
...     print(index, value)
...     index += 1
...
0 a
1 b
2 c

В этом примере index – целое число, отслеживающее, как глубоко вы находитесь в списке. На каждой итерации на экран выводится index и value. Последним шагом в цикле является обновление числа, хранящегося в index. Здесь может всплыть ошибка, если забыть обновить index:

>>> index = 0

>>> for value in values:
...     print(index, value)
...
0 a
0 b
0 c

В примере на каждой итерации index остается равным 0, т. к. нет кода для его обновления. Этот вид ошибки трудно отследить, особенно в длинных и сложных циклах.

Другой распространенный способ решения этой проблемы – использовать range() в сочетании с len() для автоматического создания index, и вам не придется помнить о его обновлении:

>>> for index in range(len(values)):
...     value = values[index]
...     print(index, value)
...
0 a
1 b
2 c

В этом примере len(values) возвращает длину value, равную 3. Затем range() создает итератор, работающий от начального value (по умолчанию 0) до тех пор, пока не будет достигнуто len(values) – 1. В цикле устанавливается value, равное элементу при текущем value index. Наконец, выводится value и index.

Одной из ошибок может быть отсутствие обновления value в начале каждой итерации.

Пример несколько ограничен, поскольку values должны разрешать доступ к своим элементам с помощью целочисленных индексов. Итерируемые объекты, разрешающие этот вид доступа, называют последовательностями.

В цикле for можно использовать любую итерацию, но только последовательности доступны по целочисленным индексам. Попытка получить доступ к элементам по индексу из генератора или итератора вызовет ошибку TypeError:

>>> enum = enumerate(values)
>>> enum[0]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'enumerate' object is not subscriptable

К счастью, функция enumerate() позволяет избежать все эти проблемы. Она доступна в Python с версии 2.3.

Использование enumerate()

Вы можете использовать enumerate() в цикле почти так же, как объект iterable. Вместо того чтобы помещать итерацию непосредственно после in в цикл for, вы помещаете ее в круглые скобки enumerate(). Придется немного изменить переменную цикла, как показано в примере:

>>> for count, value in enumerate(values):
...     print(count, value)
...
0 a
1 b
2 c

Когда вы используете enumerate(), функция возвращает две переменные цикла: количество текущих итераций и значение элемента на текущей итерации.

Как и в обычном цикле for, переменные цикла могут быть названы как угодно (например, как count и value в примере выше). Вам не нужно помнить, что следует из итератора получить доступ к элементу и подвинуть индекс в конец цикла – все делается автоматически.

Функция enumerate() имеет дополнительный аргумент, который можно использовать для управления начальным значением счетчика. По умолчанию оно равно 0, и если вы хотите получить первый элемент списка, используйте индекс 0:

>>> print(values[0])
a

Здесь можно увидеть, что доступ к значениям с индексом 0 дает первый элемент a. Бывает так, что необходимо запустить счетчик не с 0. В этом случае используйте аргумент start для enumerate(), чтобы изменить начальный счетчик:

>>> for count, value in enumerate(values, start=1):
...     print(count, value)
...
1 a
2 b
3 c

Практика с enumerate()

Допустимо применять enumerate() в любой ситуации, например, когда нужно использовать count и элемент в цикле. Имейте в виду, что функция enumerate() увеличивает количество на единицу на каждой итерации. Далее рассмотрим некоторые варианты использования enumerate().

Количество итеративных элементов

В предыдущем разделе вы видели, как использовать enumerate() для печати. Разберем скрипт, читающий файлы reST и сообщающий пользователю о проблемах с форматированием.

Не беспокойтесь, если что-то не понятно, т. к. цель – показать реальное использование enumerate():

def check_whitespace(lines):
    """Проверка наличия пробелов и проблем с длиной строки"""
    for lno, line in enumerate(lines):
        if "\r" in line:
            yield lno+1, "\\r in line"
        if "\t" in line:
            yield lno+1, "табуляция"
        if line[:-1].rstrip(" \t") != line[:-1]:
            yield lno+1, "пробел"

check_whitespace() принимает один аргумент, lines, который является строками обрабатываемого файла. Сценарий возвращает номер строки, сокращенный как lno и line, а далее идет проверка на нежелательные символы: \r, \t, табуляции и пробелы.

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

Условные операторы для пропуска элементов

Иногда есть необходимость выполнить действие на первой итерации цикла:

>>> users = ["Test User", "Real User 1", "Real User 2"]
>>> for index, user in enumerate(users):
...     if index == 0:
...         print("Вывод для:", user)
...     print(user)
...
Вывод для: Test User
Real User 1
Real User 2

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

Можно комбинировать математические операции с условиями. Например, если вам понадобится возвращать элементы из итерационной модели, имеющие четный индекс. Это делается следующим образом:

>>> def even_items(iterable):
...     """Возврат элементов ``iterable`` когда индекс четный"""
...     values = []
...     for index, value in enumerate(iterable, start=1):
...         if not index % 2:
...             values.append(value)
...     return values
...

even_items() принимает аргумент iterable. Значения инициализируются как пустой список. Затем создается цикл for с помощью enumerate() и устанавливается start=1.

В цикле for проверяется, равен ли остаток от деления index на 2 нулю и если так, то элемент добавляется к значениям.

Вы можете сделать код более понятным, используя list comprehension без инициализации пустого списка:

>>> def even_items(iterable):
...     return [v for i, v in enumerate(iterable, start=1) if not i % 2]
...

В этом примере even_items() использует list comprehension, а не цикл for для извлечения каждого элемента из списка, индекс которого является четным числом.

Вы можете проверить, что even_items() работает как ожидалось, получив четные элементы из диапазона целых чисел от 1 до 10. Результатом будет [2, 4, 6, 8, 10]:

>>> seq = list(range(1, 11))

>>> print(seq)
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

>>> even_items(seq)
[2, 4, 6, 8, 10]

even_items() возвращает четные элементы из seq. Теперь можно проверить, как эта штука будет работать с буквами алфавита ASCII:

>>> alphabet = "abcdefghijklmnopqrstuvwxyz"

>>> even_items(alphabet)
['b', 'd', 'f', 'h', 'j', 'l', 'n', 'p', 'r', 't', 'v', 'x', 'z']

Строка alphabet содержит двадцать шесть строчных букв ASCII-алфавита. Вызов even_items() и передача alphabet возвращает список чередующихся букв.

Строки в Python – это последовательности, которые можно использовать как в циклах, так и в целочисленном индексировании. В случае со строками нужны квадратные скобки для реализации той же функциональности, что и в even_items():

>>> list(alphabet[1::2])
['b', 'd', 'f', 'h', 'j', 'l', 'n', 'p', 'r', 't', 'v', 'x', 'z']

Понимание enumerate()

До этого мы рассматривали примеры использования enumerate(). Теперь стоит глубже изучить, как эта функция работает.

Чтобы лучше понять, как работает enumerate(), реализуйте собственную версию с помощью Python. Она должна следовать двум требованиям:

• принимать iterable и начальное значение в качестве аргументов;
• отправлять обратно кортеж с текущим значением счетчика и связанным с ним элементом из iterable.
  

Один из способов написания функции по данным спецификациям приведен в документации Python:

>>> def my_enumerate(sequence, start=0):
...     n = start
...     for elem in sequence:
...         yield n, elem
...         n += 1
...

my_enumerate() принимает два аргумента: sequence и start. Значение по умолчанию start равно 0. Для каждого элемента в последовательности текущие значения n и elem отправляются обратно. В конце вы увеличиваете n, чтобы подготовиться к следующей итерации. Вот my_enumerate() в действии:

>>> seasons = ["Spring", "Summer", "Fall", "Winter"]

>>> my_enumerate(seasons)
<generator object my_enumerate at 0x7f48d7a9ca50>

>>> list(my_enumerate(seasons))
[(0, 'Spring'), (1, 'Summer'), (2, 'Fall'), (3, 'Winter')]

>>> list(my_enumerate(seasons, start=1))
[(1, 'Spring'), (2, 'Summer'), (3, 'Fall'), (4, 'Winter')]

Вы создаете список четырех сезонов для работы. Далее показываете, что вызов my_enumerate() с ними в качестве последовательности создаст объект генератора. Это происходит, поскольку ключевое слово yield отправляет значения обратно вызывающему объекту.

Наконец создается два списка из my_enumerate(): в одном начальное значение остается по умолчанию 0, а в другом start изменяется на 1. В обоих случаях вы получаете список кортежей, в которых первым элементом является счетчик, а вторым – значение из seasons.

Мы реализовали эквивалент enumerate() всего из нескольких строк кода, хотя оригинальный код на C для enumerate() несколько больше. Это означает, что Python очень быстрый и эффективный.

Распаковка аргументов с помощью enumerate()

Когда вы используете enumerate() в цикле for, вы говорите Python работать с двумя переменными: одной для подсчета и одной для значения. Все это можно сделать, используя распаковку аргументов.

Идея в том, что кортеж может быть разбит на несколько переменных в зависимости от длины последовательности. Например, если распаковать кортеж из двух элементов в две переменные:

>>> tuple_2 = (10, "a")
>>> first_elem, second_elem = tuple_2
>>> first_elem
10
>>> second_elem
'a'

Сначала создается кортеж с элементами 10 и «а». Затем вы распаковываете его в first_elem и second_elem, присваивая по одному из значений кортежа.

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

>>> values = ["a", "b"]
>>> enum_instance = enumerate(values)
>>> enum_instance
<enumerate at 0x7fe75d728180>
>>> next(enum_instance)
(0, 'a')
>>> next(enum_instance)
(1, 'b')
>>> next(enum_instance)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

В этом примере создается список значений с двумя элементами «a» и «b». Затем значения передаются в enumerate() и присваивается возвращаемое значение enum_instance. Когда выводится enum_instance, можно видеть, что это экземпляр enumerate() с определенным адресом памяти.

Затем используется next(), чтобы получить следующее значение из enum_instance. Первое значение – это кортеж с числом 0 и первым элементом из значений «a».

Повторный вызов next() дает еще один кортеж, на этот раз с числом 1 и вторым элементом из значений «b». Наконец, вызов next() еще раз вызывает StopIteration, так как больше нет возвращаемых значений. Когда метод используется в цикле for, Python автоматически вызывает next() перед началом каждой итерации, пока StopIteration растет.

Если функция возвращает кортеж, можно использовать распаковку аргументов для назначения элементов кортежа нескольким переменным. Это мы делали выше, используя две переменные цикла.

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

>>> first = ["a", "b", "c"]
>>> second = ["d", "e", "f"]
>>> third = ["g", "h", "i"]
>>> for one, two, three in zip(first, second, third):
...     print(one, two, three)
...
a d g
b e h
c f i

В цикле элементы сопоставляются, а затем выводятся. Вы можете объединить zip() и enumerate(), используя вложенную распаковку аргументов:

>>> for count, (one, two, three) in enumerate(zip(first, second, third)):
...     print(count, one, two, three)
...
0 a d g
1 b e h
2 c f i

В цикле for используется вложенный zip() в enumerate(), т. е. каждый раз, когда цикл повторяется, enumerate() выдает кортеж с первым значением в качестве count и вторым в качестве другого кортежа с элементами из аргументов zip().

Существуют иные способы эмуляции поведения enumerate() в сочетании с zip(). Один из методов использует itertools.count(), который по умолчанию возвращает последовательные целые числа, начиная с нуля. Вы можете изменить предыдущий пример, чтобы это попробовать:

>>> import itertools
>>> for count, one, two, three in zip(itertools.count(), first, second, third):
...     print(count, one, two, three)
...
0 a d g
1 b e h
2 c f i

Заключение

enumerate() позволяет писать Pythonic-циклы, когда необходимо рассчитать счетчик и значение объекта interable. Большое преимущество заключается в том, что он возвращает кортеж с этими параметрами, поэтому не нужно увеличивать счетчик самостоятельно. Это также дает возможность изменить начальное значение счетчика.

В этой статье мы изучили:

• использование enumerate() в циклах for;
• применение enumerate() в реальных примерах;
  
• способы получения значения из enumerate(), используя распаковку аргументов;
  
• возможность реализации собственной функции enumerate().
  

Вы узнали много нового и теперь способны упростить ваши циклы, чтобы сделать код на Python более профессиональным.

Источники

• https://realpython.com/python-enumerate/