'''
This module provides a function that constructs a list containing
the sizes of directories under a specified directory.

If you run it as a script and want to see a list of all files >=
a specified size in MB, add a second integer or float parameter to
the command line.

Copyright (C) 2005 Don Peterson

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 (at your option) 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.

You should have received a copy of the GNU General Public
License along with this program; if not, write to the Free
Software Foundation, Inc., 59 Temple Place, Suite 330,
Boston, MA 02111-1307  USA
'''

import os
__version__ = "$Id: space.py 1.4 2005/12/24 05:42:11 donp Exp $"

# Contains directory name and total size
DirSizes = []

# Data for finding biggest files
NumBigFiles = 20   # How many to track
BigFiles    = []   # List of (size, filename)
Threshold   = 0    # MB threshold; if nonzero, show all file sizes
                   # that are >= to this value.

def TrimBigFiles():
    '''Trim the container.  If Threshold is nonzero, we keep all
    files larger than the threshold.  Otherwise, just keep the
    specified number.
    '''
    global BigFiles
    BigFiles.sort()
    if Threshold:
        BigFiles = [x for x in BigFiles if x[0] > Threshold*1e6]
    else:
        BigFiles = BigFiles[-NumBigFiles:]

def GetTotalFileSize(dummy_param, directory, list_of_files):
    '''Given a list of files and the directory they're in, add the
    total size and directory name to the global list DirSizes.
    '''
    global DirSizes, BigFiles
    currdir = os.getcwd()
    os.chdir(directory)
    total_size = 0
    if len(list_of_files) != 0:
        for file in list_of_files:
            if file == ".." or file == ".":  continue
            size = os.stat(file)[6]
            total_size = total_size + size
            BigFiles.append((size, os.path.join(directory, file)))
    DirSizes.append([total_size, directory])
    TrimBigFiles()
    os.chdir(currdir)

def GetSize(directory):
    '''Returns a list of the form [ [a, b], [c, d], ... ] where
    a, c, ... are the number of total bytes in the directory and
    b, d, ... are the directory names.  The indicated directory 
    is recursively descended and the results are sorted by directory 
    size with the largest directory at the beginning of the list.
    '''
    import os
    global DirSizes
    DirSizes = []
    os.path.walk(directory, GetTotalFileSize, "")
    DirSizes.sort()
    DirSizes.reverse()

def ShowBiggestDirectories(directory):
    import string
    GetSize(directory)
    # Get total number of bytes
    total_size = 0L
    for dir in DirSizes:
        total_size = total_size + dir[0]
    if total_size != 0:
        print "For directory '%s':    " % directory,
        print "[total space = %.1f MB]" % (total_size / 1e6)
        print "   %     MB   Directory"
        print "------ -----  " + "-" * 50
        not_shown_count = 0
        for dir in DirSizes:
            percent = 100.0 * dir[0] / total_size
            dir[1] = string.replace(dir[1], "\\\\", "/")
            if percent >= 0.1:
                print "%6.1f %5d  %s" % (percent, int(dir[0]/1e6), dir[1])
            else:
                not_shown_count = not_shown_count + 1
        if not_shown_count > 0:
            if not_shown_count > 1:
                print "            [%d directories not shown]" % not_shown_count
            else:
                print "            [%d directory not shown]" % not_shown_count

if __name__ == '__main__':
    import sys
    name = sys.argv[0]
    sys.argv = sys.argv[1:]
    if len(sys.argv) == 0:
        sys.argv.append(".")
    if len(sys.argv) == 2:
        Threshold = float(sys.argv[1])
    ShowBiggestDirectories(sys.argv[0])
    if Threshold:
        print "\nFiles >=", sys.argv[1], "MB:"
    else:
        print "\n%d biggest files in MB:" % NumBigFiles
    BigFiles.reverse()
    total = 0L
    for size, file in BigFiles:
        print "%8.2f  %s" % (size/1e6, file.replace("\\", "/"))
        total += size
    print "\n  Total MB of these files = %.1f" % (total/1e6)