#!/usr/bin/env python

# Copyright 2001 Dan York

# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or any later version.
#
# This program is distributed in the hope that it will be useful
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 
# GNU General Public License for more details.
# http://www.gnu.org/copyleft/gpl.html


# import various modules that we need

from os import *
from os.path import *
import getopt
from sys import argv,stderr,exit

# define two constants purely to make code more readable
TRUE = 1
FALSE = 0

def FormatFile (file):
    """adds the size to the filename if -s is set"""
    if Size:
       try:
          if islink(file):
	      s = lstat(file)[6]
	  else:
              s =  getsize(file)
             # format size to look prettier
          fs = ("[%9s]  " + file) % (str(s))
          return (fs)
        # if we get 'permission denied' print a bogus size field 
       except OSError:
          return ('[   denied]' + file)
    else:
       return (file)

def PrintLink(istring,file):
    """prints out a link"""
    if exists(file):
        print istring +  FormatFile(file) + " --> " + readlink(file)
    else:
        print istring + FormatFile(file) + " --> (" + readlink(file)+")"

def PrintTree(level, dir, files, prefix):
    """Prints the tree of files"""
       # alpha sort the file
    files.sort()
       # test if we are running in the cwd. If not cd to wherever we are to run
    if split(getcwd())[-1] != dir:
        chdir(dir)

       # strain out all unnecessary files - this is necessary so that we can
       # correctly determine the end of the listing for directories and therefore
       # round out the line (when a -d is used). While we are at it, let's get
       # rid of all 'hidden' files that begin with a dot.
    RevisedFileList = []
    if DirsOnly:
        for i in files:
           if i[0] != '.' and isdir(i):
              RevisedFileList.append(i)
    else:
        for i in files:
           if i[0] != '.':
              RevisedFileList.append(i)

       # loop through all the files that are left
    for f in RevisedFileList:
             # create the string to put in front of the filename
             # check if it is the last one if so, round the corner by using a quote
             # Also, set up the prefix for the next recursion (if there is one)
           if f == RevisedFileList[-1]:
               istring = prefix + "`-- "
               NextPrefix = prefix + '    '
           else:
               istring = prefix + "|-- "
               NextPrefix = prefix + '|   '
  
           try:
                   # if it is a directory, print it and recurse into that dir
               if isdir(f):
	           if islink(f):
		       PrintLink(istring,f)
		   else:
                       print istring + FormatFile(f)

		   if not islink(f):
                     PrintTree(level + 1, f, listdir(f),NextPrefix)
                   # if -d is set (dirs only) then do nothing else. otherwise proceed
               elif not DirsOnly:
                       # if it is a link, test to see if it exists
                    if islink(f):
		       PrintLink(istring,f)
                       # otherwise it must be a file
                    else:
                        print istring + FormatFile(f)
              # this is here in case a file generates "permission denied"
              # for instance if perms are wrong or the file does not exist
              # Note that directories are already printed above regardless
              # of their status
           except OSError:
               if not DirsOnly and not isdir(f):
                   print istring +  FormatFile(f)
      # change back to the directory we originally came from (only really
      # applies to recursive calls to PrintTree
    chdir('..')

# Set initial indentation level to zero, and set DirsOnly and Size to FALSE
level = 0
DirsOnly = FALSE
Size = FALSE

# Get the options and arguments and trap for errors
# Note that getopt allows for '-d -s', '-s -d', '-ds', '-sd'
# 'opts' will contain command line options, if any
# 'args' will contain other arguments, if any

try:
    opts, args = getopt.getopt(argv[1:],"ds")
except getopt.error, msg:
    # If the user gave something other than -d or -s, quit with an error
    stderr.write("\ntree: commandline error: "+msg+"\n\n")
    exit(1)

# Scan option flags and set appropriate variables
for i in opts:
    if i[0] == '-d':
        DirsOnly = TRUE
    elif i[0] == '-s':
        Size = TRUE

# if no argument, assume the current dir and print '.', else print cwd
if not args:
    cwd = getcwd()
    print "."
else:
    # check to be sure that the argument *is* a directory and if so, print it
    if isdir(args[0]):
        cwd = args[0]
        print cwd
    else:
        # if arg is NOT a directory, print error message and quit
        stderr.write("\ntree: commandline error: "+args[0]+\
	    " is not a directory\n\n")
        exit(1)

PrintTree(level,cwd,listdir(cwd),'')