FileTreeMaker

FileTreeMaker2.py

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

"""FileTreeMaker.py: ..."""

__author__  = "legendmohe"

import os
import argparse
import time

def main(args):
    rootDir = args.root
    root_level = rootDir.count(os.path.sep)
    output = args.output if len(args.output) != 0 else "output_%d" % int(time.time()) + ".txt"
    exf = args.exclude_folder
    exn = args.exclude_name
    max_layer = args.max_layer
    full_path = args.full_path
    
    print("root:%s" % rootDir)
    
    buf = []
    for dirName, subdirList, fileList in os.walk(rootDir):
        level = dirName.count(os.path.sep) - root_level
        space_prefix = "  "*level
        
        if level > max_layer and max_layer != -1:
            continue

        if full_path is True:
            buf.append(space_prefix + "-[%s]" % dirName)
        else:
            paths = os.path.split(dirName)
            buf.append(space_prefix + "-[%s]" % paths[-1])
        for folder_name in exf:
            if folder_name in subdirList:
                subdirList.remove(folder_name)
        if len(fileList) != 0:
            # buf.append(space_prefix + "  |")
            for fname in fileList:
                if not any(exclude_name in fname for exclude_name in exn):
                    buf.append(space_prefix + "  |- %s" % fname)
            buf.append(space_prefix + "  |")
    del buf[-1]  # remove last |
    
    output_str = "\n".join(buf)
    print(output_str)
    with open(output, 'w') as of:
        of.write(output_str)
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-r", "--root", help="root of file tree", default=".")
    parser.add_argument("-o", "--output", help="output file name", default="")
    parser.add_argument("-f", "--full-path", dest="full_path", action="store_true",
                        help="output full path", default=False)
    parser.add_argument("-exf", "--exclude_folder", nargs='*', help="exclude folder", default=[])
    parser.add_argument("-exn", "--exclude_name", nargs='*', help="exclude name", default=[])
    parser.add_argument("-m", "--max_layer", help="max layer",
                        type=int, default=-1)
    args = parser.parse_args()
    main(args)