git @ Cat's Eye Technologies tree / master script / tree
master

Tree @master (Download .tar.gz)

tree @masterraw · history · blame

#!/usr/bin/env python3

# SPDX-FileCopyrightText: In 2014, Chris Pressey, the original author of this work, placed it into the public domain.
# SPDX-License-Identifier: Unlicense
# For more information, please refer to <https://unlicense.org/>

import os
import shlex
import sys


def print_tree(dirname, indent, filename_format, maxwidth=75, excludes=None, all_files=False):
    dirnames = []
    filenames = []

    for filename in sorted(os.listdir(dirname)):
        if (not all_files) and filename.startswith('.'):
            continue
        fullname = os.path.join(dirname, filename)
        if os.path.islink(fullname):
            filenames.append(filename)
        elif os.path.isdir(fullname) and filename not in excludes:
            dirnames.append(filename)
        else:
            filenames.append(filename)

    if filename_format == 'line':
        file_string = ' '.join([shlex.quote(f) for f in filenames])
        if len(indent + file_string) > maxwidth:
            file_string = file_string[:(maxwidth - 4) - len(indent)] + '...'
        file_lines = [file_string]
    elif filename_format == 'block':
        file_lines = []
        file_line = ''
        for filename in filenames:
            filename = shlex.quote(filename)
            if len(file_line) + len(filename) + 1 >= maxwidth - len(indent):
                file_lines.append(file_line)
                file_line = ''
            file_line += filename + ' '
        file_lines.append(file_line)
    elif filename_format == 'list':
        file_lines = [shlex.quote(f) for f in filenames]
    elif filename_format == 'count':
        file_lines = [str(len(filenames))]
    else:
        raise NotImplementedError(filename_format)

    for subdirname in dirnames:
        print(indent + subdirname + '/')
        print_tree(
            os.path.join(dirname, subdirname), indent + '    ',
            filename_format, maxwidth=maxwidth, excludes=excludes,
            all_files=all_files
        )

    for file_line in file_lines:
        if file_line:
            print(indent + file_line)


def main(args):
    filename_format = 'block'
    maxwidth = 75
    excludes = ('venv', 'node_modules', '__pycache__', '.git')
    all_files = False
    try:
        import subprocess
        with open('/dev/tty') as tty:
            height, width = subprocess.check_output(['stty', 'size'], stdin=tty).split()
        maxwidth = int(width)
    except Exception as e:
        pass
    while args and args[0].startswith('-'):
        opt = args.pop(0)
        if opt in ('-f', '--full'):
            filename_format = 'list'
        elif opt in ('-1', '--1-line'):
            filename_format = 'line'
        elif opt in ('-c', '--count'):
            filename_format = 'count'
        elif opt in ('-w', '--max-width'):
            val = args.pop(0)
            maxwidth = int(val)
            if maxwidth <= 0:
                raise ValueError("Illegal --max-width '%s'" % val)
        elif opt in ('-x', '--exclude'):
            excludes = args.pop(0).split(',')
        elif opt in ('-a', '--all-files'):
            all_files = True
        else:
            raise ValueError("Unknown command-line option '%s'" % opt)
    if not args:
        args = ['.']
    print_tree(args[0], '', filename_format, maxwidth=maxwidth, excludes=excludes, all_files=all_files)


if __name__ == '__main__':
    main(sys.argv[1:])