# coding=utf-8
import os.path
import sys
import time

link1 = "─┬─"
link2 = "├──"
link3 = "───"
link4 = "└──"
link5 = " │ "
link6 = "   "
link7 = "│  "


class File(object):
    def __init__(self, abs_path, name, children=[]):
        self.abs_path = abs_path
        self.name = name
        self.children = children
        self.next = None

    def __repr__(self) -> str:
        return str({'name': self.name,
                    'children': self.children
                    })


def create_path(parent, file):
    return parent + os.sep + file


def sub_dirs(path):
    return [x for x in os.listdir(path) if os.path.isdir(create_path(path, x))]


def sub_dirs_abs(path):
    return [create_path(path, x) for x in os.listdir(path) if os.path.isdir(create_path(path, x))]


def dir_tree(path):
    parent = File(path, os.path.basename(path), [])
    dirs = sub_dirs_abs(path)
    for dir in dirs:
        kid = dir_tree(dir)
        parent.children.append(kid)

    for i in range(len(parent.children) - 1):
        parent.children[i].next = parent.children[i + 1]

    return parent


def space(num=0):
    return " " * num


def print_prefix():
    pass


def print_dir_tree(dir, link='', deep=0, prefix_print=''):
    print(prefix_print, end="")
    print(link, end="")
    print(dir.name)

    if dir.next:
        prefix_print += link7
    else:
        prefix_print += link6

    for i in range(len(dir.children) - 1):
        print_dir_tree(dir.children[i], link2, deep + 1, prefix_print)
    if len(dir.children) >= 1:
        print_dir_tree(dir.children[-1], link4, deep + 1, prefix_print)


def main(argv):
    base_path = argv[1] if len(argv) >= 2 else os.path.abspath('.')
    root = dir_tree(base_path)
    print_dir_tree(root)


if __name__ == '__main__':
    start = time.time()
    main(sys.argv)
    end = time.time()
    print("print dir tree cost {:.2f}s".format(end - start))

求醉的夕阳
4 声望0 粉丝