# coding: utf-8

import __builtin__
from functools import partial


# NOTE: hasnodes=operator.isSequenceType
# では文字列ノードが文字のシーケンスと解釈されて困ることがある。
# XXX: maxdepthを再帰呼び出しのlimitとして指定する場合、
# 根拠のある数のある数の方がいいかもしれない。
def tree_map(func, tree,
                  hasnodes=lambda x:isinstance(x,list),
                  getnodes=__builtin__.iter,
                  mapfunc=__builtin__.map,
                  depth=0, maxdepth=64):
    assert maxdepth > depth

    # NOTE: 再帰で呼び出しする関数を内部関数とすることで余分な引数引渡しを排除。
    def rec(xs, depth):
        assert maxdepth >= depth
        if hasnodes(xs):
            return mapfunc(partial(rec,depth=depth+1), getnodes(xs))
        else:
            return func(xs)

    return rec(tree, depth=depth)

# NOTE: 大抵の再帰関数はスタックとループで置き換え可能。
# NOTE: 関数呼び出しのスタックを消費しないのでmaximum recursion limitの制限なし、
# maxdepth は任意の深度まで探索したい時に指定できる。
def tree_walk(tree,
              hasnodes=lambda x:isinstance(x,list),
              getnodes=__builtin__.iter,
              with_parent=True,
              skip_deep_nodes=True,
              depth=0, maxdepth=-1,
              ):
    stack = [(tree,getnodes(tree))]
    while stack:
        parent,top = stack[-1]
        try:
            item = top.next()
        except StopIteration:
            stack.pop()
            depth -= 1
            assert len(stack) == depth+1
        else:
            if 0 <= maxdepth <= depth: # NOTE: maxdepth=-1 always False
                if skip_deep_nodes:
                    continue
            else:
                if hasnodes(item):
                    stack.append((item, getnodes(item)))
                    depth += 1
                    assert len(stack) == depth+1
                    continue
            yield (parent,item) if with_parent else item


if __name__ == '__main__':
    
    tree = [[1,2],[3,[4,5,6,[7,8,9]]]]

    print tree_map(lambda x:x*2, tree)
    # ==> [[2, 4], [6, [8, 10, 12, [14, 16, 18]]]]

    print list(tree_walk(tree, with_parent=False, maxdepth=3))
    # ==> [1, 2, 3, 4 ,5 , 6]

    # tuple用の関数をpartialを使って生成。
    tuple_tree_map = partial(tree_map,
                             hasnodes=lambda x:isinstance(x,tuple),
                             mapfunc=lambda *args: tuple(map(*args)))
    print tuple_tree_map((2).__mul__, ((1,2),(3,(4,5))))
    # ((2, 4), (6, (8, 10)))
    
    xs = [
        [[1,2,3],[4,5,6]],
        [[7,8,9]],
    ]
    for x in tree_walk(xs, skip_deep_nodes=False, with_parent=False, maxdepth=1):
        print x
    # [1,2,3]
    # [4,5,6]
    # [7,8,9]


    import os
    from os.path import join as pathjoin
    dir_walk = partial(tree_walk,
                     hasnodes=os.path.isdir,
                     getnodes=lambda path: (pathjoin(path,x) for x in os.listdir(path)))

    ## not test on codepad.org
    if 0:
        for dirname,filepath in dir_walk(u"./"):
            print filepath

[[2, 4], [6, [8, 10, 12, [14, 16, 18]]]]
[1, 2, 3, 4, 5, 6]
((2, 4), (6, (8, 10)))
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]