unchain.py

unchain.py

import sys
import itertools as it

PY2 = sys.version_info.major == 2

if PY2:
    range = xrange


def unchain(elems, chain_size):
    """Reverse of itertools.chain

    >>> unchain([1, 2, 3, 4, 5, 6], 3)
    [[1, 2, 3], [4, 5, 6]]
    """
    for i in range(0, len(elems), chain_size):
        yield elems[i:i + chain_size]


test_chained_vals = [1, 2, 3, 4, 5, 6, 7, 8]

test_unchained_vals = [
    [1, 2, 3, 4],
    [5, 6, 7, 8],
]

assert list(it.chain(*test_unchained_vals)) == test_chained_vals
assert list(unchain(test_chained_vals, 4)) == test_unchained_vals