Deep iterator from http://www.mitbbs.com/article_t/JobHunting/32583985.html

deep_iter.py

#!/usr/bin/python

class DeepIterator(object):
    def __init__(self, ls):
        """ls: A list containing integers or list of same type."""
        self.ls = ls
        # An index stack indicating current index in each level.
        self.curr = [0]
    
    def __iter__(self):
        return self
    
    def next(self):
        if self.curr[0] == len(self.ls):
            raise StopIteration
        
        inner_ls = self.ls
        for i in self.curr[:-1]:
            inner_ls = inner_ls[i]
        
        if self.curr[-1] == len(inner_ls):
            self.curr.pop()
            self.curr[-1] += 1
            return self.next()
        
        curr_ele = inner_ls[self.curr[-1]]
        if isinstance(curr_ele, int):
            self.curr[-1] += 1
            return curr_ele
        elif isinstance(curr_ele, list):
            self.curr.append(0)
            return self.next()
        else:
            raise TypeError


if __name__ == '__main__':
    ls = []
    assert list(DeepIterator(ls)) == []
    
    ls = [1, 2, 3]
    assert list(DeepIterator(ls)) == [1, 2, 3]
    
    ls = [[10, 11], [4, 5, [7, [9, [1, ]]]], 3]
    assert list(DeepIterator(ls)) == [10, 11, 4, 5, 7 ,9, 1, 3]

    ls = [[4, 5, [7, [9, [1, ]]]], [[99, 100], 88]]
    assert list(DeepIterator(ls)) == [4, 5, 7 ,9, 1, 99, 100, 88]
    
    ls = [1, [2, 9.9], 3]
    try:
        list(DeepIterator(ls))
    except TypeError:
        pass
    except Exception:
        raise AssertionError("Wrong exception")
    else:
        raise AssertionError("No exception")