pretty_tree.py

"""
Utilities for printing out a nice looking file tree, e.g

pretty_tree(('root', [
    ('a', [
        'file1.txt',
        'file2.txt',
    ]),
    ('b', [
        'file1.json',
        'file2.json',
    ])
]))

root/
├── a/
│   ├── file1.txt
│   └── file2.txt
└── b/
    ├── file1.json
    └── file2.json
"""


def pretty_tree(d):
    return '\n'.join(tree_lines(d))


def tree_lines(d, prefix='', last=False):

    def format_node(node, postfix=''):
        if len(prefix) >= 3:
            if last:
                return '{}└── {}{}'.format(prefix[:-4], node, postfix)
            else:
                return '{}├── {}{}'.format(prefix[:-4], node, postfix)
        else:
            return '{}{}'.format(node, postfix)

    if isinstance(d, str):
        return [format_node(d)]
    assert isinstance(d, tuple), '{} should be tuple'.format(d)
    dir_name, children = d

    result = [
        format_node(dir_name, postfix='/')
    ]
    last_i = len(children) - 1
    for i, child in enumerate(children):
        is_last = (i == last_i)
        child_prefix = '{}    '.format(prefix) if is_last else '{}│   '.format(prefix)
        result.extend(tree_lines(child, child_prefix, last=is_last))
    return result