#! /usr/bin/python

import os, sys, stat, getopt, re, traceback, string

__version__ = "0.3"

def showVersion():
  print 'rimdu v'+__version__+', by Rimon Barr:',
  print 'Disk usage utility.'

def showUsage():
  print
  showVersion()
  print
  print 'Usage: rimdu -e [dir]'
  print '       rimdu -s [-n=#] [-d | -f] [dir]'
  print '       rimdu -c [-n=#] [-d | -f] [dir]'
  print '       rimdu -h | -v'
  print 
  print '  -h, -?, --help      display this help information'
  print '  -v, --version       display version'
  print '  -e, --empty         show empty directories below dir'
  print '  -s, --size          show files/directories in order of size'
  print '  -c, --count         show directories in order of file count'
  print '  -t, --total         show total count/size'
  print '  -n=                 restrict listing to n entries'
  print '  -f                  files only'
  print '  -d                  directories only'
  print
  print 'Send comments, suggestions and bug reports to <barr+rimdu@cs.cornell.edu>.'
  print

def usageError():
  print 'rimdu: command syntax error'
  print 'Try `rimdu --help\' for more information.'

# format file size (bytes)
def formatBytes(bytes):
  if bytes<=9999: return "%ld" % bytes
  bytes = bytes / 1024
  if bytes<=9999: return "%ldK" % bytes
  bytes = bytes / 1024
  if bytes<=9999: return "%ldM" % bytes
  bytes = bytes / 1024
  return "%ldG" % bytes

def dircountlist(path):
  try:
    st=os.lstat(path)
    if stat.S_ISDIR(st[stat.ST_MODE]):
      cumlist = []
      cumcount = 0
      for f in os.listdir(path):
        count, list = dircountlist(os.path.join(path, f))
        if count!=-1:
          cumcount = cumcount + count
          cumlist = cumlist + list
      cumlist = [(cumcount, path)] + cumlist
      return cumcount+1, cumlist
    else:
      return 1, []
  except KeyboardInterrupt:
    raise
  except OSError:
    return -1, []
  except Exception:
    raise

def dirsizelist(path, fileonly, dironly):
  try:
    st=os.lstat(path)
    if stat.S_ISDIR(st[stat.ST_MODE]):
      cumlist=[]
      cumsize=0
      for f in os.listdir(path):
        size, list = dirsizelist(os.path.join(path, f), fileonly, dironly)
        if size!=-1:
          cumsize = cumsize + size
          cumlist = cumlist + list
      if not fileonly:
        cumlist = [(cumsize, path)] + cumlist
      return cumsize, cumlist
    else:
      if not dironly:
        return long(st[stat.ST_SIZE]), [(st[stat.ST_SIZE], path)]
      else:
        return long(st[stat.ST_SIZE]), []
  except KeyboardInterrupt: 
    raise
  except OSError:
    return -1, []
  except Exception:
    raise

def printemptydir(arg, path, sub):
  if not len(os.listdir(path)):
    print path

def rimdu():
  # parse options
  try:
    opts, args = getopt.getopt(sys.argv[1:], 'hv?esn:fdct',
      ['help', 'version', 'empty', 'size', 'count', 'total',])
  except getopt.GetoptError: 
    usageError()
    return

  # rimdu modes
  emptymode=0
  sizemode=0
  countmode=0
  cutoff=0
  fileonly=0
  dironly=0
  totalmode=0

  # process options
  for o, a in opts:
    if o in ("-h", "--help", "-?"):
      showUsage()
      return
    if o in ("-v", "--version"):
      showVersion()
      return
    if o in ('-e', '--empty'):
      emptymode=1
    if o in ('-s', '--size'):
      sizemode=1
    if o in ('-c', '--count'):
      countmode=1
    if o == '-n':
      cutoff=int(re.match('=(\d+)', a).group(1))
    if o == '-f':
      fileonly=1
    if o == '-d':
      dironly=1
    if o in ('-t', '--total'):
      totalmode=1
  if fileonly and dironly:
    usageError()
    return
  if len(args)==1 and args[0]=='-':
    paths = sys.stdin.readlines()
    paths = map(string.strip, paths)
  elif args: paths=args
  else: paths=['.',]

  total = 0
  for path in paths:
    # operate
    if emptymode:
      os.path.walk(path, printemptydir, None)
    elif sizemode:
      _, list = dirsizelist(path, fileonly, dironly)
      list.sort()
      list.reverse()
      if cutoff:
        list=list[:cutoff]
      for size, f in list:
        print "%6s" % formatBytes(size), f
        total = total + size
    elif countmode:
      _, list = dircountlist(path)
      list.sort()
      list.reverse()
      if cutoff:
        list=list[:cutoff]
      for count, f in list:
        print "%6ld" % count, f
        total = total + count
    else:
      usageError()
      return
  if totalmode:
    if sizemode:
      print "%6s" % formatBytes(total)
    if countmode:
      print "%6ld" % total

try:
  if __name__=='__main__': rimdu()
except KeyboardInterrupt:
  traceback.print_exc()
  print 'Break!'