Iterate an iterator by chunks (of n) in Python?

Можете ли вы придумать хороший способ (возможно, с помощью itertools) разбить итератор на части заданного размера?

Следовательно, l=[1,2,3,4,5,6,7] with chunks(l,3) становится итератором [1,2,3], [4,5,6], [7]

Я могу придумать небольшую программу для этого, но не лучший способ, возможно, с помощью itertools.

grouper() Рецепт из itertoolsрецептов документации близок к тому, что вы хотите:

def grouper(iterable, n, *, incomplete='fill', fillvalue=None):
    "Collect data into non-overlapping fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3, fillvalue='x') --> ABC DEF Gxx
    # grouper('ABCDEFG', 3, incomplete='strict') --> ABC DEF ValueError
    # grouper('ABCDEFG', 3, incomplete='ignore') --> ABC DEF
    args = [iter(iterable)] * n
    if incomplete == 'fill':
        return zip_longest(*args, fillvalue=fillvalue)
    if incomplete == 'strict':
        return zip(*args, strict=True)
    if incomplete == 'ignore':
        return zip(*args)
    else:
        raise ValueError('Expected fill, strict, or ignore')

Однако это не будет хорошо работать, когда последний фрагмент будет неполным, поскольку, в зависимости от incomplete режима, он либо заполнит последний фрагмент значением fill, вызовет исключение, либо автоматически удалит неполный фрагмент.

В более поздние версии рецептов они добавили batched рецепт, который делает именно то, что вы хотите:

def batched(iterable, n):
    "Batch data into tuples of length n. The last batch may be shorter."
    # batched('ABCDEFG', 3) --> ABC DEF G
    if n < 1:
        raise ValueError('n must be at least one')
    it = iter(iterable)
    while (batch := tuple(islice(it, n))):
        yield batch

Наконец, менее общее решение, которое работает только с последовательностями, но обрабатывает последний фрагмент по желанию и сохраняет тип исходной последовательности, - это:

(my_list[i:i + chunk_size] for i in range(0, len(my_list), chunk_size))

Хотя OP просит функцию возвращать фрагменты в виде списка или кортежа, на случай, если вам нужно вернуть итераторы, то решение Свена Марнаха может быть изменено:

def batched_it(iterable, n):
    "Batch data into iterators of length n. The last batch may be shorter."
    # batched('ABCDEFG', 3) --> ABC DEF G
    if n < 1:
        raise ValueError('n must be at least one')
    it = iter(iterable)
    while True:
        chunk_it = itertools.islice(it, n)
        try:
            first_el = next(chunk_it)
        except StopIteration:
            return
        yield itertools.chain((first_el,), chunk_it)

Некоторые тесты: http://pastebin.com/YkKFvm8b

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

python 3.12 добавляет itertools.пакетированный, который работает со всеми итерабельными объектами (включая списки):

>>> from itertools import batched
>>> for batch in batched('ABCDEFG', 3):
...     print(batch)
('A', 'B', 'C')
('D', 'E', 'F')
('G',)

Начиная с python 3.8, существует более простое решение, использующее оператор :=:

def grouper(iterator: Iterator, n: int) -> Iterator[list]:
    while chunk := list(itertools.islice(iterator, n)):
        yield chunk

а затем вызвать его таким образом:

>>> list(grouper(iter('ABCDEFG'), 3))
[['A', 'B', 'C'], ['D', 'E', 'F'], ['G']]

Примечание: вы можете поместить iter в grouper функцию, чтобы принимать Iterable вместо Iterator .