iter_chunks - Split a python iterable in chunks, loading one at a time.

iter_chunks.py

import itertools


def iter_chunks(iterable, chunk_size):
    """Split the iterable in chunks of the given size, loading one chunk at a time.
    This function yields lists of size between 1 and `chunk_size`.

    :param iterable: the iterable to iterate over
    :type iterable: typing.Iterable[T]
    :param chunk_size: the size to use for each chunk.
    :type chunk_size: int

    :rtype: typing.Generator[typing.List[T], None, None]
    """
    if chunk_size <= 0:
        raise ValueError("Chunk size must be a positive integer")

    it = iter(iterable)
    while True:
        chunk = list(itertools.islice(it, chunk_size))
        if not chunk:
            break
        yield chunk

        
        
###### Tests #######
import pytest

@pytest.mark.parametrize(
    'iterable, chunk_size, expected',
    [
        (range(0), 3, []),
        (range(11), 3, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10]]),
        (range(10), 3, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]),
        (range(9), 3, [[0, 1, 2], [3, 4, 5], [6, 7, 8]]),
    ]
)
def test_iter_chunk(iterable, chunk_size, expected):
    results = []
    for chunk in iter_chunks(iterable, chunk_size=chunk_size):
        results.append(chunk)

    assert results == expected


@pytest.mark.parametrize('chunk_size', [0, -1, -10, None])
def test_iter_chunk_invalid_size(chunk_size):
    with pytest.raises((ValueError, TypeError)):
        next(iter_chunks(range(20), chunk_size=chunk_size))