Wraps an iterator so that StopIteration is only raised once

reuse_guard.py

#!/usr/bin/env python3.8
from typing import TypeVar, Iterator

"""
https://docs.python.org/3/library/stdtypes.html#iterator-types
This is considered broken by the iterator protocol, however I think
that what's considered broken is to continue to _yield values_, where
with this we emphasize the fact that if ``StopIteration`` is raised
once, the iterator _should not be used_ further. Those are two
different things.
"""

class IteratorExhaustedError(Exception):
    """Exception raised when exhausted iterators are ``next``d"""

T = TypeVar("T")
class reuse_guard(Iterator[T]):
    """
    Wraps an iterator so that StopIteration is only raised once,
    after that, ``IteratorExhaustedError`` will be raised to detect
    fixed-size iterator misuses
    """
    def __init__(self, iterator: Iterator[T]):
        self._iterated: bool = False
        self._iterator = iterator

    def __next__(self) -> T:
        try:
            return next(self._iterator)
        except StopIteration as e:
            if self._iterated:
                raise IteratorExhaustedError(
                    "This iterator has already reached its end")
            self._iterated = True
            raise e
    
    def __iter__(self) -> Iterator[T]:
        return self