strict_generator.py

#
# THE PROBLEM
#

# create a generator and list out its items
a = (n**2 for n in range(10))
print(list(a))

# listing it out again creates an empty list because the generator has been consumed;
# generators in this state will continue to raise StopIteration, but caller will take
# this to mean an empty generator
print(list(a))

#
# DEFINE A StrictGenerator WRAPPER
# that raises an exception if a consumed generator gets accessed again.
#

class AccessingConsumedGeneratorException(Exception):
    """
    Exception raised if a generator is accessed after it has already raised StopIteration.
    """

class StrictGenerator:
    """
    Wrapper for generator which will only permit it to be consumed once.
    Additional accesses will raise AccessingConsumedGeneratorException.
    """
    def __init__(self, gen):
        self._gen = gen
        self._gen_ended = False

    def __iter__(self):
        return self

    def __next__(self):
        try:
            ret = next(self._gen)
            return ret
        except StopIteration:
            if self._gen_ended:
                ex = AccessingConsumedGeneratorException()
                raise ex from None
            self._gen_ended = True
            raise

# wrap a generator in StrictGenerator
a = StrictGenerator(n**2 for n in range(10))

# consume a as usual
print(list(a))

# subsequent access of a will raise AccessingConsumedGeneratorException
print(list(a))