Как не пропустить следующий элемент после itertools.takewhile()

Скажем, мы хотим обработать итератор и хотим обрабатывать его кусками.
Логика на кусок зависит от ранее вычисленных кусков, поэтому groupby() не помогает.

Проблема заключается в том, что takewhile() нужно пройти мимо последнего элемента, который соответствует новой логике блоков, таким образом "съедая" первый элемент для следующего фрагмента.

Существуют различные решения, включая упаковку или à la C ungetc() и т.д.

Мой вопрос: есть ли элегантное решение?

Ответы

Ответ 1

takewhile() действительно нужно посмотреть на следующий элемент, чтобы определить, когда переключать поведение.

Вы можете использовать обертку, которая отслеживает последний замеченный элемент, и которая может быть "reset" для резервного копирования одного элемента:

_sentinel = object()

class OneStepBuffered(object):
    def __init__(self, it):
        self._it = iter(it)
        self._last = _sentinel
        self._next = _sentinel
    def __iter__(self):
        return self
    def __next__(self):
        if self._next is not _sentinel:
            next_val, self._next = self._next, _sentinel
            return next_val
        try:
            self._last = next(self._it)
            return self._last
        except StopIteration:
            self._last = self._next = _sentinel
            raise
    next = __next__  # Python 2 compatibility
    def step_back(self):
        if self._last is _sentinel:
            raise ValueError("Can't back up a step")
        self._next, self._last = self._last, _sentinel

Оберните свой итератор в этом, прежде чем использовать его с помощью takewhile():

myIterator = OneStepBuffered(myIterator)
while True:
    chunk = itertools.takewhile(getNewChunkLogic(), myIterator)
    process(chunk)
    myIterator.step_back()

Демо:

>>> from itertools import takewhile
>>> test_list = range(10)
>>> iterator = OneStepBuffered(test_list)
>>> list(takewhile(lambda i: i < 5, iterator))
[0, 1, 2, 3, 4]
>>> iterator.step_back()
>>> list(iterator)
[5, 6, 7, 8, 9]

Ответ 2

Учитывая, что вызываемый GetNewChunkLogic() будет сообщать True вдоль первого фрагмента и False после.
Следующий фрагмент

решает проблему "дополнительного следующего шага" takewhile.
является элегантным, потому что вам не нужно реализовывать логику "назад-один шаг".

def partition(pred, iterable):
    'Use a predicate to partition entries into true entries and false entries'
    # partition(is_odd, range(10)) -->  1 3 5 7 9 and 0 2 4 6 8
    t1, t2 = tee(iterable)
    return filter(pred, t1), filterfalse(pred, t2)

while True:
    head, tail = partition(GetNewChunkLogic(), myIterator)
    process(head)
    myIterator = tail

Однако наиболее элегантным способом является преобразование вашего GetNewChunkLogic в генератор и удаление цикла while.