aboutsummaryrefslogtreecommitdiff
path: root/yapymake/util/peekable_iterator.py
blob: 67bd59d10975c5834ffb3f07147c5af4470211d9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from typing import Iterator, Optional, TypeVar

__all__ = [
    'PeekableIterator',
]

T = TypeVar('T')

class PeekableIterator(Iterator[T]):
    _inner: Iterator[T]
    _peeked: Optional[T]

    def __init__(self, inner: Iterator[T]):
        self._inner = inner
        self._peeked = None

    def __iter__(self) -> Iterator[T]:
        return self

    def __next__(self) -> T:
        if self._peeked is not None:
            result = self._peeked
            self._peeked = None
            return result
        return next(self._inner)

    def peek(self) -> T:
        if self._peeked is None:
            self._peeked = next(self._inner)
        return self._peeked