| from collections.abc import Iterable, Iterator | |
| class tracked_str(str): | |
| origins = {} | |
| def set_origin(self, origin: str): | |
| if super().__repr__() not in self.origins: | |
| self.origins[super().__repr__()] = origin | |
| def get_origin(self): | |
| return self.origins.get(super().__repr__(), str(self)) | |
| def __repr__(self) -> str: | |
| if super().__repr__() not in self.origins or self.origins[super().__repr__()] == self: | |
| return super().__repr__() | |
| else: | |
| return f"{str(self)} (origin={self.origins[super().__repr__()]})" | |
| class tracked_list(list): | |
| def __init__(self, *args, **kwargs) -> None: | |
| super().__init__(*args, **kwargs) | |
| self.last_item = None | |
| def __iter__(self) -> Iterator: | |
| for x in super().__iter__(): | |
| 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})" | |
| 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)) | |