unzip.py

unzip.py

import itertools


def unzip(in_: Iterable[Tuple[Any, ...]]) -> Iterable[Tuple[Any, ...]]:
    """Unzips a N-sized iterable of M-tuples to a M-sized iterable of N-tuples.

    Example:
        >>> data = [('a', 'b', 'c'), (1, 2, 3)]  # 2x3
        >>> unzip(data) == [('a', 1), ('b', 2), ('c', 3)]  # 3x2
        True
    """

    in_iter = iter(in_)
    in_first, in_rest = next(in_iter), in_iter

    out = zip(*itertools.chain((in_first,), in_rest))

    try:
        # Only return the iterable if there are values.
        out_first, out_rest = next(out), out
        return itertools.chain((out_first,), out_rest)
    except StopIteration:
        # Make sure there's always an M-sized iterable as output for uniformity.
        # Without this logic, `zip()` would result in an empty iterable.
        return itertools.repeat((), len(in_first))