Python BufferedIterator

buffered_iterator.py

"""
buffered_iterator - Decorates an iterator with "has_next" by buffering values.

Example use case for this is if you have a base interface/ABC that has methods 
`has_next` and `get_next`, and several classes realizing this interface already 
would have iterators internally, you wouldn't want them to each have to 
implement the logic for `has_next` - so it would be better to decorate an 
iterator with some sort of `has_next`.

Inspiration from 
https://stackoverflow.com/questions/1966591/hasnext-in-python-iterators
"""


# Copyright (c) 2022, Tim Kutcher
# 
# MIT License - See https://github.com/tkutcher/tkutcher/blob/master/LICENSE.md


import typing
from collections.abc import Iterator

_T = typing.TypeVar("_T")
_T_co = typing.TypeVar("_T_co", covariant=True)


_EMPTY_BUF = object()


class BufferedIterator(Iterator[_T]):
    def __init__(self, real_it: Iterator[_T]):
        self._real_it = real_it
        self._buf = next(self._real_it, _EMPTY_BUF)

    def has_next(self):
        return self._buf is not _EMPTY_BUF

    def __next__(self) -> _T_co:
        v = self._buf
        self._buf = next(self._real_it, _EMPTY_BUF)
        if v is _EMPTY_BUF:
            raise StopIteration()
        return v

test_buffered_iterator.py

# Copyright (c) 2022, Tim Kutcher
# 
# MIT License - See https://github.com/tkutcher/tkutcher/blob/master/LICENSE.md


import pytest

from buffered_iterator import BufferedIterator


class TestBufferedIterator:
    def test_can_use_like_normal(self):
        it = BufferedIterator(iter([1, 2, 3]))
        assert next(it) == 1
        assert next(it) == 2
        assert next(it) == 3
        with pytest.raises(StopIteration):
            next(it)

        for v in BufferedIterator(iter([1, 1, 1])):
            assert v == 1

    def test_can_use_has_next(self):
        it = BufferedIterator(iter([1, 2, 3]))
        assert it.has_next()
        assert next(it) == 1
        assert next(it) == 2
        assert it.has_next()
        assert next(it) == 3
        assert not it.has_next()

    def test_empty_iterator(self):
        it = BufferedIterator(iter([]))
        assert not it.has_next()

    def test_can_iterate_past_none(self):
        it = BufferedIterator(iter([1, None, 2]))
        assert next(it) == 1
        assert it.has_next()
        assert next(it) is None
        assert it.has_next()
        assert next(it) == 2