YOOOO

flatten_iterator.py

"""
Receieves an iterator of iterators, implement flattened iterator
"""

class FlattenIterator:
    def __init__(self, iterator):
        self.iterator = iterator
 
    def __iter__(self):
        # Handle empty
        try:
            self.current = next(self.iterator)
        except StopIteration:
            self.current = iter([])
        return self

    def __next__(self):
        while True:
            try:
                return next(self.current)
            except StopIteration:
                # We allow this to raise StopIteration -> means that we've handled all
                self.current = next(self.iterator)


# Test
def test_flatten_iterator(test_arrays, expected_output):
    iterator = iter(iter(arr) for arr in test_arrays)
    f = FlattenIterator(iterator=iterator)
    res = [i for i in f]
    assert res == expected_output

arrays = [
    [1, 2, 3],
    [],
    [4, 5],
    []
]
expected = [1, 2, 3, 4, 5]
test_flatten_iterator(arrays, expected)

arrays = []
expected = []
test_flatten_iterator(arrays, expected)