在Python中每n项拆分生成器/可迭代对象（splitEvery）

Question

我正在尝试用Python编写一个叫做'splitEvery'的Haskell函数。它的定义如下：

splitEvery :: Int -> [e] -> [[e]]
    @'splitEvery' n@ splits a list into length-n pieces.  The last
    piece will be shorter if @n@ does not evenly divide the length of
    the list.

这个基本版本运行得很好，但我想要一个可以处理生成器表达式、列表和迭代器的版本。而且，如果输入是一个生成器，那么输出也应该是一个生成器！

测试

# should not enter infinite loop with generators or lists
splitEvery(itertools.count(), 10)
splitEvery(range(1000), 10)

# last piece must be shorter if n does not evenly divide
assert splitEvery(5, range(9)) == [[0, 1, 2, 3, 4], [5, 6, 7, 8]]

# should give same correct results with generators
tmp = itertools.islice(itertools.count(), 10)
assert list(splitEvery(5, tmp)) == [[0, 1, 2, 3, 4], [5, 6, 7, 8]]

当前实现

这是我现在的代码，但它在处理简单列表时不太好用。

def splitEvery_1(n, iterable):
    res = list(itertools.islice(iterable, n))
    while len(res) != 0:
        yield res
        res = list(itertools.islice(iterable, n))

这个版本在处理生成器表达式时也不行（感谢jellybean帮我修复）：

def splitEvery_2(n, iterable): 
    return [iterable[i:i+n] for i in range(0, len(iterable), n)]

一定有一段简单的代码可以做到分割。我知道我可以写不同的函数，但这似乎应该是个简单的事情。我可能在一个不重要的问题上卡住了，但这真的让我很烦。

---

这个功能类似于来自http://docs.python.org/library/itertools.html#itertools.groupby的grouper，但我不想让它填充额外的值。

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

它确实提到了一种截断最后一个值的方法。这也不是我想要的。

迭代器的从左到右的评估顺序是有保证的。这使得使用izip(*[iter(s)]*n)将数据系列聚类成n长度组成为可能。

list(izip(*[iter(range(9))]*5)) == [[0, 1, 2, 3, 4]]
# should be [[0, 1, 2, 3, 4], [5, 6, 7, 8]]

Answer 1

more_itertools 是一个库，它里面有一个叫做 chunked 的功能：

import more_itertools as mit


list(mit.chunked(range(9), 5))
# [[0, 1, 2, 3, 4], [5, 6, 7, 8]]

Answer 2

这里有一个简洁的一行代码版本。和Haskell一样，它是懒惰的。

from itertools import islice, takewhile, repeat
split_every = (lambda n, it:
    takewhile(bool, (list(islice(it, n)) for _ in repeat(None))))

这段代码需要你在调用 split_every 之前先使用 iter。

示例：

list(split_every(5, iter(xrange(9))))
[[0, 1, 2, 3, 4], [5, 6, 7, 8]]

虽然下面的版本不是一行代码，但它不需要你调用 iter，这可以避免一个常见的错误。

from itertools import islice, takewhile, repeat

def split_every(n, iterable):
    """
    Slice an iterable into chunks of n elements
    :type n: int
    :type iterable: Iterable
    :rtype: Iterator
    """
    iterator = iter(iterable)
    return takewhile(bool, (list(islice(iterator, n)) for _ in repeat(None)))

（感谢@eli-korvigo的改进。）

Answer 3

from itertools import islice

def split_every(n, iterable):
    i = iter(iterable)
    piece = list(islice(i, n))
    while piece:
        yield piece
        piece = list(islice(i, n))

>>> list(split_every(5, range(9)))
[[0, 1, 2, 3, 4], [5, 6, 7, 8]]

>>> list(split_every(3, (x**2 for x in range(20))))
[[0, 1, 4], [9, 16, 25], [36, 49, 64], [81, 100, 121], [144, 169, 196], [225, 256, 289], [324, 361]]

>>> [''.join(s) for s in split_every(6, 'Hello world')]
['Hello ', 'world']

>>> list(split_every(100, []))
[]

一些测试：