3

After reading How do you split a list into evenly sized chunks? and seeing this kind of mistakes happeing https://code.djangoproject.com/ticket/18972 all the time.

Why is not a chunk function in itertools?

Edit: grouper from http://docs.python.org/2/library/itertools.html#recipes doesn't have the same behavior as chunks

Example:

chunks([1, 2, 3, 4, 5], 3)
# Should return [[1, 2, 3], [4, 5]] or the iterator equivalent.
Community
  • 1
  • 1
razpeitia
  • 1,947
  • 4
  • 16
  • 36

2 Answers2

4

Posting a question like this here is not the way to get something like this added to Python. You should maybe try a Python mailing list.

I implemented chunks() for you with the semantics you requested. Getting the handling of the last chunk just right was a little bit tricky, but otherwise this is pretty easy. If it were added to itertools it would be written in C so it would be faster.

Tested and working in Python 2.6, Python 2.7, and Python 3.2.

import itertools as it
import sys

# use lazy xrange on 2.x; on 3.x plain "range" is always lazy
if sys.version_info[0] < 3:
    _range = xrange
else:
    _range = range

def chunks(iterable, n):
    """
    Yield up lists of n elements, taken from iterable.
    If length of iterable is not evenly divisible by n, the last list will be short.
    """
    if n < 1:
        raise ValueError("n must be >= 1")

    itr = iter(iterable)
    try:
        while True:
            lst = []
            for _ in _range(n):
                lst.append(next(itr))
            if not lst:
                break
            yield lst
    except StopIteration:
        # Only yield up a partial chunk if it is not zero length.
        if lst:
            yield lst

print(list(chunks([1, 2, 3, 4, 5, 6], 3)))  # prints: [[1, 2, 3], [4, 5, 6]]
print(list(chunks([1, 2, 3, 4, 5], 3))) # prints: [[1, 2, 3], [4, 5]]
print(list(chunks([], 3))) # prints: []
print(list(chunks([1, 2], 0))) # raises ValueError exception

EDIT:

The inefficiency of the above solution was kinda bugging me. I was pretty sure there had to be a simpler solution using itertools.islice() so I figured it out. I like this one much better.

def chunks(iterable, n):
    """
    Yield up lists of n elements, taken from iterable.
    If length of iterable is not evenly divisible by n, the last list will be short.
    """
    if n < 1:
        raise ValueError("n must be >= 1")

    itr = iter(iterable)

    while True:
        lst = list(it.islice(itr, n))
        if not lst:
            break
        yield lst
Eric O. Lebigot
  • 91,433
  • 48
  • 218
  • 260
steveha
  • 74,789
  • 21
  • 92
  • 117
2

It's not in itertools, but it's mentioned on the very page for itertools as a recipe:

http://docs.python.org/2/library/itertools.html#recipes

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)

You may as well ask why all the other recipes aren't included in itertools :)

Patashu
  • 21,443
  • 3
  • 45
  • 53