Разбираемся в работе функции grouper в Python

Часто на StackOverflow задают вопрос «Как разбить последовательность на части равного размера?». Давайте попробуем в нём разобраться.

Если это действительно последовательность, а не произвольный итерируемый объект, то разбить её можно с помощью срезов. Однако в документации itertools есть отличный способ сделать это с помощью функции grouper.

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

Пары объектов

Начнём с более простой функции для группировки итератора по объектам чётной длины в итератор по парам объектов:

			def pairs(iterable):
       it = iter(iterable)
       return zip(it, it)
		

Как это работает?

Сначала мы создаём итератор по итерируемому объекту. Итератор — объект, который может выдавать элементы по одному при вызове функции next(). Итератор сам по себе является итерируемым объектом: при вызове iter() он возвращает сам себя. Чаще всего мы встречаем итераторы в виде объектов, которые возвращаются генераторами, функциями map, filter, zip, функциями в itertools и т.д. Однако мы можем создать итератор для любого итерируемого объекта с помощью функции iter. Например:

			>>> a = range(5) # не является итератором
>>> list(a)
[0, 1, 2, 3, 4]
>>> list(a)
[0, 1, 2, 3, 4]
>>> i = iter(a) # является итератором
>>> list(i)
[0, 1, 2, 3, 4]
>>> list(i)
[]
		

Поскольку мы прошлись по i во время первого вызова, во время второго там уже ничего нет. Для лучшего понимания давайте взглянем на функции вроде islice или takewhile, которые поглощают только часть итератора:

			>>> from itertools import islice
>>> i = iter(a)
>>> list(islice(i, 3))
[0, 1, 2]
>>> list(islice(i, 3))
[3, 4]
		

Возможно вы задаётесь вопросом, что бы случилось, будь a итератором изначально. В этом случае iter просто возвращает a.

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

			>>> a = range(10)
>>> i1, i2 = iter(a), iter(a) # два отдельных итератора
>>> list(i1)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>>> list(i2)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>>> i1 = i2 = iter(a) # две ссылки на один и тот же итератор
>>> list(i1)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>>> list(i2)
[]
		

Само собой, если a уже являлся итератором, вызов iter(a) вернул бы нам сам а дважды, и результаты выполнения первого и второго примеров были бы одинаковы.

А что произойдет, если применить функцию zip к двум ссылкам на один итератор? Каждая получит через одно значение:

			>>> i1 = i2 = iter(a)
>>> list(zip(i1, i2))
[(0, 1), (2, 3), (4, 5), (6, 7), (8, 9)]
		

Если вам сложно понять, почему так происходит, посмотрите, как работает этот упрощенный zip на чистом Python:

			def fake_zip(i1, i2):
     while True:
         v1 = next(i1)
         v2 = next(i2)
         yield v1, v2
		

Если i1 и i2 — один и тот же итератор, после строки v1 = next(i1) i1 и i2 будут указывать на следующее значение после v1, поэтому  v2 = next(i2) получит это значение.

Прим. перев. Этот бесконечный цикл завершится с исключением StopIteration, когда один из итераторов закончится. Но начиная с версии Python 3.7, выброшенное исключение StopIteration внутри генератора не приведёт к его нормальному завершению, а будет перевыброшено снаружи как RuntimeError. Таким образом, для поддержки новых версий Python этот код нужно модифицировать. Подробнее можно почитать на официальном сайте Python.

Вот и всё, что нужно знать о функции pairs.

Части произвольного размера

А как же сделать n ссылок на один и тот же итератор? Для этого есть несколько способов, но вот самый простой:

args = [iter(iterable)] * n

Теперь нужно придумать, как использовать zip в этом случае. Так как zip принимает любое количество аргументов, а не ограничивается двумя, мы можем использовать распаковку списка аргументов:

zip(*args)

И теперь мы почти можем написать grouper:

			def grouper(iterable, n):
     args = [iter(iterable)] * n
     return zip(*args)
		

Части разного размера

И наконец, что будет, если количество элементов нельзя поделить на части равного размера? Например, что произойдет, если вы захотите получить группы по 3 элемента для range(10)? Вот несколько возможных вариантов:

1. [(0, 1, 2), (3, 4, 5), (6, 7, 8), (9, None, None)]
2. [(0, 1, 2), (3, 4, 5), (6, 7, 8), (9, 0, 0)]
3. [(0, 1, 2), (3, 4, 5), (6, 7, 8), (9,)]
4. [(0, 1, 2), (3, 4, 5), (6, 7, 8)]
5. Будет выброшено исключение ValueError

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

Для достижения этой цели в библиотеке itertools есть функция zip_longest (izip_longest в Python 2.x). Вместо отсутствующих значений она вставляет None или аргумент fillvalue. Например:

			>>> list(zip_longest(*iters, fillvalue=0))
[(0, 1, 2), (3, 4, 5), (6, 7, 8), (9, 0, 0)]
		

И вот у нас есть всё, что нужно, чтобы написать grouper и понимать, как он работает:

			def grouper(iterable, n, fillvalue=None):
       # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
       args = [iter(iterable)] * n
       return zip_longest(*args, fillvalue=fillvalue)