python:是否有用于分块输入流的库函数？

Question

我想分 block 输入流以进行批处理。给定一个输入列表或生成器，

x_in = [1, 2, 3, 4, 5, 6 ...]

我想要一个函数来返回该输入的 block 。比如说，如果 chunk_size=4，那么，

x_chunked = [[1, 2, 3, 4], [5, 6, ...], ...]

这是我反复做的事情，想知道是否有比自己编写更标准的方法。我在 itertools 中遗漏了什么吗？ (可以使用 enumerate 和 groupby 解决问题，但感觉很笨拙。)如果有人想看实现，就在这里，

def chunk_input_stream(input_stream, chunk_size):
    """partition a generator in a streaming fashion"""
    assert chunk_size >= 1
    accumulator = []
    for x in input_stream:
        accumulator.append(x)
        if len(accumulator) == chunk_size:
            yield accumulator
            accumulator = []
    if accumulator:
        yield accumulator

编辑

受 kreativitea 的回答启发，这里有一个使用 islice 的解决方案，它很简单并且不需要后过滤，

from itertools import islice

def chunk_input_stream(input_stream, chunk_size):
    while True:
        chunk = list(islice(input_stream, chunk_size))
        if chunk:
            yield chunk
        else:
            return

# test it with list(chunk_input_stream(iter([1, 2, 3, 4]), 3))

Answer 1

来自 itertools 的配方:

def grouper(n, iterable, fillvalue=None):
    "Collect data into fixed-length chunks or blocks"
    # grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx
    args = [iter(iterable)] * n
    return izip_longest(fillvalue=fillvalue, *args)