Super Chunkify

super_custom_chunkify.py

# -*- coding: utf-8 -*-
"""
На работе как-то раз появилась довольно стандартная задача:
    написать генаратор, который отдаёт элементы из некого итератора пачками
    по n элементов.
Дополнительное условие:
    Не возвращать пустой список на последней итерации (возникает, когда
    общее количество элементов кратно размеру пачки).

В общем, как-то её в тот раз решили, хоть и не слишком красиво.

А сегодня захотелось ещё одного дополнительного условия:
    На первой итерации отдавать пачку размером n - 1 элементов.

Всю дорогу домой в голове крутились итераторы, генераторы, чанкифайеры... :)
Решил выплеснуть всё в код, чтобы спокойно спать.
"""

# Сначала в голову пришёл такой лаконичный вариант простого chunkify:
def super_chunkify(iterable, n):
    """Отдаёт списки по n элементов из iterable."""

    iterator = iter(iterable)

    buf = [iterator.next()]     # В буфере держим один запасной элемент.

    while True:
        try:
            for _ in range(n):  # Добавляем в буфер n элементов - имеем n+1.
                buf.append(iterator.next())
        except StopIteration:   # Если получили m < n элементов,
            yield buf           # то отдаём весь буфер, т.е m <= n элементов
            raise StopIteration # и выходим.
        else:
            yield buf[:n]       # А если хватило - отдаём ровно n элементов,
            del buf[:n]         # и удаляем их из буфера.


"""
В первый раз написал c finally.
Получилось сильно короче, но выглядит менее очевидно и выбрасывает последний
буфер при любом эксепшене, что плохо:

        try:
            for _ in range(n):
                buf.append(iterator.next())
        finally:
            yield buf[:n]
            del buf[:n]
"""



# Потом решил сделать размер chunk'ов динамическим с помощью send():
def super_custom_chunkify(iterable, n):
    """Отдаёт списки по n элементов, пока в генератор не передано новое n."""

    iterator = iter(iterable)

    buf = [iterator.next()]       # В буфере держим один запасной элемент.

    while True:
        try:
            for _ in range(n):    # Добавляем в буфер n элементов - имеем n+1.
                buf.append(iterator.next())
        except StopIteration:     # Если получили m < n элементов,
            yield buf             # то отдаём весь буфер, т.е m <= n элементов
            raise StopIteration   # и выходим.
        else:
            new_n = yield buf[:n] # А если хватило - отдаём ровно n элементов,
            del buf[:n]           # и удаляем их из буфера,
            if new_n:             # и сохраняем новое n, пришедшее из send'а.
                n = new_n


# Ну и решил твою задачу с помощью super_custom_chunkify():
def special_chunkify(iterable, n):
    """Отдаёт списки из n-1, n, n, ... элементов."""
    # Настраиваем отдачу по n-1 элементу.
    scc = super_custom_chunkify(iterable, n - 1)
    yield scc.next()      # Забираем и отдаём n-1 элемент.
    yield scc.send(n)     # Настраиваем отдачу по n, забираем и отдаём n.
    while True:
        yield scc.next()  # Отдаём по n элементов.


# print list(special_chunkify([], 3))
# print list(special_chunkify(range(7), 3))
# print list(special_chunkify(range(8), 3))
# print list(special_chunkify(range(9), 3))

# []
# [0, 1], [2, 3, 4], [5, 6]]
# [[0, 1], [2, 3, 4], [5, 6, 7]]
# [[0, 1], [2, 3, 4], [5, 6, 7], [8]]