tree.py

#!/usr/bin/env python3

from pathlib import Path
import functools
import operator
import os.path
import stat
import sys


BLOCK_SIZES = (
    (1024*1024*1024*1024, 'T'),
    (1024*1024*1024,      'G'),
    (1024*1024,           'M'),
    (1024,                'K'),
)

def format_size(size):
    for bs, unit in BLOCK_SIZES:
        if size >= bs:
            return '%.1f%s' % (size / bs, unit)
    return str(size)


def tree(path, dirs_only=False, max_depth=0, _depth=0):
    path = Path(path)
    lst = path.lstat()
    is_symlink = stat.S_ISLNK(lst.st_mode)
    st = lst if is_symlink else path.stat()
    is_dir = stat.S_ISDIR(st.st_mode)
    if is_symlink:
        size = 0
    elif is_dir:
        size = functools.reduce(operator.add, [
            tree(p, dirs_only=dirs_only,
                 max_depth=max_depth,
                 _depth=_depth+1)
            for p in sorted(path.iterdir())
        ], 0)
    else:
        size = lst.st_size
    if (is_dir or not dirs_only) and \
       (not max_depth or _depth <= max_depth):
        p = str(path)
        if is_dir:
            p += os.path.sep
        if is_symlink:
            p += ' -> ' + os.readlink(str(path))
        print('%10s  %s' % (format_size(size), p))
    return size


def main():
    args = []
    max_depth = 0
    dirs_only = False
    argv = iter(sys.argv[1:])
    for a in argv:
        if a == '-d':
            dirs_only = True
            continue
        if a == '-L':
            max_depth = int(next(argv))
            continue
        if a.startswith('-'):
            raise ValueError(a)
        args.append(a)
    for a in args:
        tree(a, dirs_only=dirs_only, max_depth=max_depth)


if __name__ == '__main__':
    main()