Eclectic Media Git ptree / ee349a3 tree / __init__.py
ee349a3

Tree @ee349a3 (Download .tar.gz)

__init__.py @ee349a3raw · history · blame

#! /usr/bin/env python

# tree.py
#
# Written by Doug Dahms
#
# Prints the tree structure for the path specified on the command line

from os import listdir, sep
from os.path import abspath, basename, isdir

from argparse import ArgumentParser


def tree(origin_path, padding, print_files=False, truncate_files=False, truncate_dirs=False):
    """ :truncate_dirs: If true, don't follow dir depths deeper than 3. """
    print(padding[:-1] + '+-' + basename(abspath(origin_path)) + '/')
    padding = padding + ' '
    files = []

    dirs = [x for x in listdir(origin_path) if isdir(origin_path + sep + x)]
    files = [x for x in listdir(origin_path) if not isdir(origin_path + sep + x)]

    if print_files:
        padding += '| '

        for i, file in enumerate(files):
            print(padding + '+-' + file)

            if len(files) > 3 and i + 1 == 3 and truncate_files:
                print(padding + '+- ...')
                break

        padding = padding[:-2]

    # count = 0
    for i, dir in enumerate(dirs):

        path = origin_path + sep + dir

        if truncate_dirs:
            if len(dirs) < 3:
                tree(path, padding + '|', print_files, truncate_files, truncate_dirs)
            else:
                print(padding + '+-' + basename(abspath(dir)) + '/')

            if i + 1 == 3:
                print(padding + '+- ...')
                break

        else:
            tree(path, padding + '|', print_files, truncate_files, truncate_dirs)


def _main():
    parser = ArgumentParser()

    parser.add_argument('path', default='./', type=str,
                        help='The path to map.')

    parser.add_argument('-f', '--print_files', action='store_true',
                        help='Output files alongside directories.')

    parser.add_argument('-F', '--truncate_files', action='store_true',
                        help='Only output 3 files of any given level. '
                        '(Implies `-f`)')

    parser.add_argument('-D', '--truncate_directories', action='store_true',
                        help='Only output 3 directories of any given level.')

    parsed = parser.parse_args()

    if parsed.truncate_files:
        parsed.print_files = True

    if isdir(parsed.path):
        tree(parsed.path, ' ', parsed.print_files, parsed.truncate_files, parsed.truncate_directories)
    else:
        print('%s is not a valid path...' % parsed.path)


if __name__ == '__main__':
    _main()