class TrackedIterableFromGenerator(Iterable):
    """Utility class to create an iterable from a generator function, in order to reset the generator when needed."""

    def __init__(self, generator, *args):
        super().__init__()
        self.generator = generator
        self.args = args
        self.last_item = None

    def __iter__(self):
        for x in self.generator(*self.args):
            self.last_item = x
            yield x
        self.last_item = None

    def __repr__(self) -> str:
        if self.last_item is None:
            return super().__repr__()
        else:
            return f"{self.__class__.__name__}(current={self.last_item})"

    def __reduce__(self):
        return (self.__class__, (self.generator, *self.args))
/Users/nielsrogge/Documents/python_projecten/datasets/src/datasets/utils/track.py