在Python中按块（每块n个）迭代迭代器？

Question

你能想到一个好方法（也许可以用itertools库）来把一个迭代器分成指定大小的小块吗？

比如说，给定一个列表 l=[1,2,3,4,5,6,7]，如果用 chunks(l,3) 来处理，它会变成一个迭代器，结果是 [1,2,3], [4,5,6], [7]。

我能想到一个简单的程序来实现这个功能，但我想不出一个更好的方法，也许可以用itertools来做。

Answer 1

Python 3.12 新增了一个叫做 itertools.batched 的功能，它可以作用于所有可迭代的对象（包括列表）：

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

Answer 2

虽然提问者希望函数返回的是列表或元组，但如果你需要返回迭代器的话，可以对Sven Marnach的解决方案进行一些修改：

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

只有当你的函数在每个块中遍历元素时，这种方法才会稍微高效一些。

Answer 3

来自 itertools 文档的 grouper() 这个方法，差不多能满足你的需求：

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 模式的不同，它可能会用一个填充值来填充最后一组，抛出一个异常，或者默默地丢掉这组不完整的数据。

在最近的版本中，他们添加了 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))

从 Python 3.12 开始，你也可以直接使用 itertools.batched。在 文档中这样写：

itertools.batched(iterable, n)

将可迭代对象中的数据分批成长度为 n 的元组。最后一批可能会短于 n。