https://leetcode.com/problems/peeking-iterator/

Python

class PeekingIterator:
    def __init__(self, iterator):
        """
        Initialize your data structure here.
        :type iterator: Iterator
        """
        self.data = iterator
        self.peeked = None

    def peek(self):
        """
        Returns the next element in the iteration without advancing the iterator.
        :rtype: int
        """
        if self.peeked is not None:
            return self.peeked

        if self.data.hasNext():
            self.peeked = self.data.next()
            return self.peeked

    def next(self):
        """
        :rtype: int
        """
        if self.peeked is not None:
            result = self.peeked
            self.peeked = None
            return result
        
        if self.data.hasNext():
            return self.data.next()

    def hasNext(self):
        """
        :rtype: bool
        """
        return self.data.hasNext() if self.peeked is None else True