Solution: Chunked Iterator, Python Riddle 
Tags: python
Published: 2019-08-03 22:59:18

This is a detailed solution to [this riddle](https://dev.to/orenovadia/chunked-iterator-python-riddle-1md3).

Here it is again:  write a function (`chunks`) where the input is an iterator. And the output is an iterator of `n` sized iterators.


The first challenge is that the length of the original iterator is unknown. Let's start with a naive broken solution using `itertools.islice` to create `n` size iterators without caring about the length of the original iterator:

```python
import itertools
from typing import TypeVar, Iterator

T = TypeVar('T')


def chunks(iterator: Iterator[T], n: int) -> Iterator[Iterator[T]]:
    while True:
        yield itertools.islice(iterator, 0, n)  # yield `n` size chunk of the iterator


for chunk in chunks(iter(range(10)), 3):
    print(f'A {type(chunk)} type chunk, contains items: {tuple(chunk)}')
```
The output is:
```
A <class 'itertools.islice'> type chunk, contains items: (0, 1, 2)
A <class 'itertools.islice'> type chunk, contains items: (3, 4, 5)
A <class 'itertools.islice'> type chunk, contains items: (6, 7, 8)
A <class 'itertools.islice'> type chunk, contains items: (9,)
A <class 'itertools.islice'> type chunk, contains items: ()
A <class 'itertools.islice'> type chunk, contains items: ()
.... forever

```

Using `itertools.islice` we managed to chunk up the original iterator, but we don't know when it is exhausted. And so, `chunks` is a generator function that never ends. 

To overcome this problem we need to take one item out of the original iterator. Then, we'll use `itertools.chain` to create a chunk featuring this one item and `n-1` more items.

```python

def chunks(iterator: Iterator[T], n: int) -> Iterator[Iterator[T]]:
    for first in iterator: # take one item out (exits loop if `iterator` is empty)
        rest_of_chunk = itertools.islice(iterator, 0, n - 1)
        yield itertools.chain([first], rest_of_chunk)  # concatenate the first item back

```

The output is:
```
A <class 'itertools.chain'> type chunk, contains items: (0, 1, 2)
A <class 'itertools.chain'> type chunk, contains items: (3, 4, 5)
A <class 'itertools.chain'> type chunk, contains items: (6, 7, 8)
A <class 'itertools.chain'> type chunk, contains items: (9,)
```

And That is it!

And in comprehension form (I usually prefer comprehensions, but I must admit that this one is not as readable as the loop form):
```python
    return (itertools.chain([first], itertools.islice(iterator, 0, n - 1))
            for first in iterator)
```

### Performance 
The runtime overhead of `islice`-ing and `chain`-ing is way smaller than writing your own custom code for `chunks` since both these functions are implemented in `C` (for the `CPython` runtime).

### The Caveat
There is a caveat here: This whole solution assumes that the consumer of `chunks` is consuming the iterators fully and in order. If that is not the case, the order of items in our chunks might not be consistent with the original iterator, due to the laziness of `chunks`.

Example:
```python 
chunk_iterator = chunks(iter(range(10)), 3)
first_chunk = next(chunk_iterator)
second_chunk = next(chunk_iterator)
print(tuple(second_chunk))  # (1, 2, 3)
print(tuple(first_chunk))  # (0, 4, 5)
```

Cheers!