zip iterators asserting for equal length in python

An optional boolean keyword argument, strict, is introduced for the built-in zip function in PEP 618.

Quoting What’s New In Python 3.10:

The zip() function now has an optional strict flag, used to require that all the iterables have an equal length.

When enabled, a ValueError is raised if one of the arguments is exhausted before the others.

>>> list(zip('ab', range(3)))[('a', 0), ('b', 1)]>>> list(zip('ab', range(3), strict=True))Traceback (most recent call last):  File "<stdin>", line 1, in <module>ValueError: zip() argument 2 is longer than argument 1

I can think of a simpler solution, use itertools.zip_longest() and raise an exception if the sentinel value used to pad out shorter iterables is present in the tuple produced:

from itertools import zip_longestdef zip_equal(*iterables):    sentinel = object()    for combo in zip_longest(*iterables, fillvalue=sentinel):        if sentinel in combo:            raise ValueError('Iterables have different lengths')        yield combo

Unfortunately, we can't use zip() with yield from to avoid a Python-code loop with a test each iteration; once the shortest iterator runs out, zip() would advance all preceding iterators and thus swallow the evidence if there is but one extra item in those.

Here is an approach that doesn't require doing any extra checks with each loop of the iteration. This could be desirable especially for long iterables.

The idea is to pad each iterable with a "value" at the end that raises an exception when reached, and then do the needed verification only at the very end. The approach uses zip() and itertools.chain().

The code below was written for Python 3.5.

import itertoolsclass ExhaustedError(Exception):    def __init__(self, index):        """The index is the 0-based index of the exhausted iterable."""        self.index = indexdef raising_iter(i):    """Return an iterator that raises an ExhaustedError."""    raise ExhaustedError(i)    yielddef terminate_iter(i, iterable):    """Return an iterator that raises an ExhaustedError at the end."""    return itertools.chain(iterable, raising_iter(i))def zip_equal(*iterables):    iterators = [terminate_iter(*args) for args in enumerate(iterables)]    try:        yield from zip(*iterators)    except ExhaustedError as exc:        index = exc.index        if index > 0:            raise RuntimeError('iterable {} exhausted first'.format(index)) from None        # Check that all other iterators are also exhausted.        for i, iterator in enumerate(iterators[1:], start=1):            try:                next(iterator)            except ExhaustedError:                pass            else:                raise RuntimeError('iterable {} is longer'.format(i)) from None

Below is what it looks like being used.

>>> list(zip_equal([1, 2], [3, 4], [5, 6]))[(1, 3, 5), (2, 4, 6)]>>> list(zip_equal([1, 2], [3], [4]))RuntimeError: iterable 1 exhausted first>>> list(zip_equal([1], [2, 3], [4]))RuntimeError: iterable 1 is longer>>> list(zip_equal([1], [2], [3, 4]))RuntimeError: iterable 2 is longer