Zip-итераторы, утверждающие равную длину в python

Я ищу хороший способ zip нескольких итераций, создающих исключение, если длины итераций не равны.

В случае, когда итераторы являются списками или имеют метод len, это решение является чистым и легким:

def zip_equal(it1, it2):
    if len(it1) != len(it2):
        raise ValueError("Lengths of iterables are different")
    return zip(it1, it2)

Однако, если it1 и it2 являются генераторами, предыдущая функция выходит из строя, потому что длина не определена TypeError: object of type 'generator' has no len().

Я предполагаю, что модуль itertools предлагает простой способ реализовать это, но пока я не смог его найти. Я придумал это домашнее решение:

def zip_equal(it1, it2):
    exhausted = False
    while True:
        try:
            el1 = next(it1)
            if exhausted: # in a previous iteration it2 was exhausted but it1 still has elements
                raise ValueError("it1 and it2 have different lengths")
        except StopIteration:
            exhausted = True
            # it2 must be exhausted too.
        try:
            el2 = next(it2)
            # here it2 is not exhausted.
            if exhausted:  # it1 was exhausted => raise
                raise ValueError("it1 and it2 have different lengths")
        except StopIteration:
            # here it2 is exhausted
            if not exhausted:
                # but it1 was not exhausted => raise
                raise ValueError("it1 and it2 have different lengths")
            exhausted = True
        if not exhausted:
            yield (el1, el2)
        else:
            return

Решение можно протестировать с помощью следующего кода:

it1 = (x for x in ['a', 'b', 'c'])  # it1 has length 3
it2 = (x for x in [0, 1, 2, 3])     # it2 has length 4
list(zip_equal(it1, it2))           # len(it1) < len(it2) => raise
it1 = (x for x in ['a', 'b', 'c'])  # it1 has length 3
it2 = (x for x in [0, 1, 2, 3])     # it2 has length 4
list(zip_equal(it2, it1))           # len(it2) > len(it1) => raise
it1 = (x for x in ['a', 'b', 'c', 'd'])  # it1 has length 4
it2 = (x for x in [0, 1, 2, 3])          # it2 has length 4
list(zip_equal(it1, it2))                # like zip (or izip in python2)

Я не вижу альтернативного решения? Есть ли более простая реализация моей функции zip_equal?

PS: Я написал вопрос мышления в Python 3, но также приветствуется решение Python 2.

Ответы

Ответ 1

Я могу упростить решение, используйте itertools.zip_longest() и создаю исключение, если значение отправителя, используемое для вырезания более коротких итераций, присутствует в полученном кортеже:

from itertools import zip_longest

def 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

К сожалению, мы не можем использовать zip() с yield from, чтобы избежать цикла Python-кода с тестом на каждую итерацию; после истечения кратчайшего итератора zip() будет продвигать все предыдущие итераторы и, таким образом, проглатывает доказательства, если в них есть только один дополнительный элемент.

Ответ 2

Вот подход, который не требует каких-либо дополнительных проверок с каждым циклом итерации. Это может быть желательно особенно для длинных итераций.

Идея состоит в том, чтобы поместить каждый итерабельный с "значением" в конце, который вызывает исключение при достижении, а затем делает необходимую проверку только в самом конце. Подход использует zip() и itertools.chain().

Код ниже был написан для Python 3.5.

import itertools

class ExhaustedError(Exception):
    def __init__(self, index):
        """The index is the 0-based index of the exhausted iterable."""
        self.index = index

def raising_iter(i):
    """Return an iterator that raises an ExhaustedError."""
    raise ExhaustedError(i)
    yield

def 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

Ниже показано, как он выглядит.

>>> 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

Ответ 3

Я придумал решение, использующее invirel iterable FYI:

class _SentinelException(Exception):
    def __iter__(self):
        raise _SentinelException


def zip_equal(iterable1, iterable2):
    i1 = iter(itertools.chain(iterable1, _SentinelException()))
    i2 = iter(iterable2)
    try:
        while True:
            yield (next(i1), next(i2))
    except _SentinelException:  # i1 reaches end
        try:
            next(i2)  # check whether i2 reaches end
        except StopIteration:
            pass
        else:
            raise ValueError('the second iterable is longer than the first one')
    except StopIteration: # i2 reaches end, as next(i1) has already been called, i1 length is bigger than i2
        raise ValueError('the first iterable is longger the second one.')