Second One

flatten_iterator.py

"""
Create an iterator that yields elements of nested list.
"""


class FlatIterator:
    def __init__(self, data):
        self.stack = [iter(data)]

    def __iter__(self):
        return self

    def __next__(self):
        if not self.stack:
            raise StopIteration

        current = self.stack.pop()
        try:
            element = next(current)
            self.stack.append(current)
            
            if isinstance(element, list):
                self.stack.append(iter(element))
                # If we've found a list, 
                # we need to iterate it instead
                return next(self)
            else:
                return element
        except StopIteration:
            # If current iterator has finished
            # We need to try next one
            return next(self)


a = [
	[1, 2, 3],
	[4, [5], 6],
	[],
	[[7]]
]

flat = FlatIterator(a)

assert [i for i in flat] == list(range(1, 8))