diff options
| -rw-r--r-- | iter_utils.py | 42 |
1 files changed, 42 insertions, 0 deletions
diff --git a/iter_utils.py b/iter_utils.py index 977cf1d..b1a137c 100644 --- a/iter_utils.py +++ b/iter_utils.py @@ -54,6 +54,48 @@ class PeekingIterator(Iterator): return None +class PushbackIterator(Iterator): + """An iterator that allows you to push items back + onto the front of the sequence. e.g. + + >>> i = PushbackIterator(iter(range(3))) + >>> i.__next__() + 0 + >>> i.push_back(99) + >>> i.push_back(98) + >>> i.__next__() + 98 + >>> i.__next__() + 99 + >>> i.__next__() + 1 + >>> i.__next__() + 2 + >>> i.push_back(100) + >>> i.__next__() + 100 + >>> i.__next__() + Traceback (most recent call last): + ... + StopIteration + """ + + def __init__(self, source_iter: Iterator): + self.source_iter = source_iter + self.pushed_back: List[Any] = [] + + def __iter__(self) -> Iterator: + return self + + def __next__(self) -> Any: + if len(self.pushed_back): + return self.pushed_back.pop() + return self.source_iter.__next__() + + def push_back(self, item: Any): + self.pushed_back.append(item) + + class SamplingIterator(Iterator): """An iterator that simply echoes what source_iter produces but also collects a random sample (of size sample_size) of the stream that can |