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

Tree @master (Download .tar.gz)

__init__.py @masterraw · history · blame

#! /usr/bin/env python
""" Prints the tree structure for the path specified on the command line

Original written by Doug Dahms (see http://code.activestate.com/recipes/217212-treepy-graphically-displays-the-directory-structur/)

Current version written by Ariana Giroux <me@arianagiroux.ca>, circa 2019
"""

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

from argparse import ArgumentParser


def tree(origin_path, padding, print_files=False, exclude_path=None, truncate_files=False, truncate_dirs=False):
    """ Traverses the file system, outputting file structure as it goes.

    :origin_path: The path to start the traversal from.
    :print_files: If true, output files in directories.
    :truncate_files: If true, don't output more than 3 files per level.
    :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)]
    if exclude_path is not None:
        for path in exclude_path:
            if dirs.count(path):
                dirs.remove(path)

    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]

    for i, dir in enumerate(dirs):

        path = origin_path + sep + dir

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

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

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


def _main():
    """ Main function, runs tree function.

    Also ensures clean run time, performs basic sanitation on command line
    arguments."""

    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('-e', '--exclude_paths', type=str,
                        help='A comma separated list of paths to exclude from the'
                        'traversal. (Where \'.git,__pycache__\' yields a list with '
                        '\'.git\' and \'__pycache__\' as separate paths)')

    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 parsed.exclude_paths is not None:
        exclude_path = parsed.exclude_paths.split(',')
    else:
        exclude_path = None

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


if __name__ == '__main__':
    _main()