python实现目录树生成示例

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import optparse

location_none     = 'none'
location_mid      = 'mid'
location_mid_gap  = 'mid_gap'
location_tail     = 'tail'
location_tail_gap = 'tail_gap'

notations = {
    location_none: '',
    location_mid: '├─',
    location_mid_gap: '│  ',
    location_tail: '└─',
    location_tail_gap: '    '
}

class node(object):
    def __init__(self, name, depth, parent=none, location=location_none):
        self.name = name
        self.depth = depth
        self.parent = parent
        self.location = location
        self.children = []

    def __str__(self):
        sections = [self.name]
        parent = self.has_parent()
        if parent:
            if self.is_tail():
                sections.insert(0, notations[location_tail])
            else:
                sections.insert(0, notations[location_mid])
            self.__insert_gaps(self, sections)
        return ''.join(sections)

    def __insert_gaps(self, node, sections):
        parent = node.has_parent()
        # parent exists and parent's parent is not the root node
        if parent and parent.has_parent():
            if parent.is_tail():
                sections.insert(0, notations[location_tail_gap])
            else:
                sections.insert(0, notations[location_mid_gap])
            self.__insert_gaps(parent, sections)

    def has_parent(self):
        return self.parent

    def has_children(self):
        return self.children

    def add_child(self, node):
        self.children.append(node)

    def is_tail(self):
        return self.location == location_tail

class tree(object):
    def __init__(self):
        self.nodes = []

    def debug_print(self):
        for node in self.nodes:
            print(str(node) + '/')

    def write2file(self, filename):
        try:
            with open(filename, 'w') as fp:
                fp.writelines(str(node) + '/\n'
                              for node in self.nodes)
        except ioerror as e:
            print(e)
            return 0
        return 1

    def build(self, path):
        self.__build(path, 0, none, location_none)

    def __build(self, path, depth, parent, location):
        if os.path.isdir(path):
            name = os.path.basename(path)
            node = node(name, depth, parent, location)
            self.add_node(node)
            if parent:
                parent.add_child(node)

            entries = self.list_folder(path)
            end_index = len(entries) - 1
            for i, entry in enumerate(entries):
                childpath = os.path.join(path, entry)
                location = location_tail if i == end_index else location_mid
                self.__build(childpath, depth + 1, node, location)

    def list_folder(self, path):
        """folders only."""
        return [d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))]
        # for entry in os.listdir(path):
        #     childpath = os.path.join(path, entry)
        #     if os.path.isdir(childpath):
        #         yield entry

    def add_node(self, node):
        self.nodes.append(node)

def _parse_args():
    parser = optparse.optionparser()
    parser.add_option(
        '-p', '--path', dest='path', action='store', type='string',
        default='./', help='the path to generate the tree [default: %default]')
    parser.add_option(
        '-o', '--out', dest='file', action='store', type='string',
        help='the file to save the result [default: pathname.trees]')
    options, args = parser.parse_args()
    # positional arguments are ignored
    return options

def main():
    options = _parse_args()
    path = options.path
    if not os.path.isdir(path):
        print('%s is not a directory' % path)
        return 2

    if not path or path == './':
        filepath = os.path.realpath(__file__)  # for linux
        path = os.path.dirname(filepath)
    tree = tree()
    tree.build(path)
    # tree.debug_print()
    if options.file:
        filename = options.file
    else:
        name = os.path.basename(path)
        filename = '%s.trees' % name
    return tree.write2file(filename)

if __name__ == '__main__':
    import sys
    sys.exit(main())

运行效果

gtest_start/
├─build/
├─lib/
│  └─gtest/
├─output/
│  ├─primer/
│  │  ├─debug/
│  │  │  ├─lib/
│  │  │  └─obj/
│  │  └─release/
│  │      ├─lib/
│  │      └─obj/
│  └─thoughts/
│      ├─debug/
│      │  ├─lib/
│      │  └─obj/
│      └─release/
│          ├─lib/
│          └─obj/
├─src/
│  ├─primer/
│  └─thoughts/
├─test/
│  ├─primer/
│  └─thoughts/
├─third_party/
│  └─gtest/
└─tools/